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

arXiv:2610.17132 · cs.LG · Submitted 2026-10-12 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Masked Attention Pooling Is Invariant to Padding and Permutation".

Jane: The gist The paper defines and formally proves that masked attention pooling is invariant to padding and permutation using Lean 4, connecting these mathematical invariances to a JAX kernel implementation.

Tom: First, who's behind it and why it matters.

Paper summary: Tom: So, we're looking at this paper right now, "Masked Attention Pooling Is Invariant to Padding and Permutation: A Formal Proof in Lean." The core idea is that when you take a variable-size bag of instance embeddings—like tiles from a whole slide image—and pool them into one vector using attention, that final pooled output shouldn't change if you just pad the bag with extra rows or shuffle the order of those rows around.

Jane: That sounds like something we deal with all the time in machine learning when we have batched data where some samples might be shorter than others. So, what’s the main claim here about this paper?

Tom: Well, it formally proves that masked attention pooling is invariant to padding and permutation using Lean four and mathlib <ref:2610.17132#pg1,masked attention pooling is invariant to padding and permutation>. It claims this isn't just an assumption; they built a rigorous mathematical proof that shows these properties hold exactly for the defined construction <ref:2610.17132#pg5>.

Lu: What's interesting is how they connect those abstract mathematical invariances to a JAX kernel implementation. They trace a JAX kernel, specifically `masked pool`, at shapes like three two and five two, showing the certified theorems hold for those concrete shapes <ref:2610.17132#pg6>.

Meng: From a production standpoint, that’s good because we worry about what breaks when we scale things up. If you have padding or reordering messes up your results, it creates instability in the model output. So proving invariance means the pooling layer itself is robust to those common data handling issues <ref:2610.17132#pg8>.

Lalam: I see what they are doing here. They define a guarded normalizer `D` to fix the value on empty bags and avoid division by zero, and then they prove that if no row is valid or the bag is empty, the output is zero <ref:2610.17132#pg6>.

Jane: So, what does that guarding do in practice? Does it solve the problem of division by zero when we have an empty bag of valid instances?

Tom: Exactly. They show that if no row is valid, or the index set `I` is empty, the numerator becomes a zero vector and the final output becomes zero <ref:2610.17132#pg6>. That’s a key part of their proof for padding invariance <ref:2610.17132#pg5>.

Paper summary: Lu: They use size-generic Lean four theorems to prove things like invariance under any bijection of the row index type and under appending any finite number of masked rows with arbitrary real features and logits <ref:2610.17132#pg5>. That’s a very broad statement about the structure of their definition.

Meng: Broad is good, but I want to know how this translates to what we actually run on our hardware. Does this Lean proof mean that when I deploy this kernel, I don't need to worry about my specific padding strategy as long as it’s finite?

Tom: The paper confirms that the certified theorems hold for arbitrary finite index types and feature dimensions <ref:2610.17132#pg5>. And they verified the JAX kernel meets these specifications on CPU devices <ref:2610.17132#pg6>.

Lalam: I also see a corollary that applies to instance-local scorers, which states that for any function `f` and permutation `e`, the pooling operation behaves predictably across those reorderings <ref:2610.17132#pg6>.

Jane: That sounds like it means the relationship between the input features and the final output is stable regardless of how we organize those inputs <ref:2610.17132#pg5>. So, what’s the bigger picture for someone who just uses this AI tool without reading Lean?

Tom: This work is focused strictly on proving structural properties of the aggregation layer itself and doesn't make claims about accuracy or clinical validity <ref:2610.17132#pg9>. They are concerned with the mathematical consistency of the pooling mechanism, not whether it’s a good classifier.

Lu: It really zeroes in on proving that the way attention is averaged over valid instances is stable under changes to the indexing structure <ref:2610.17132#pg5>. It shows that if you define your masking and normalization correctly, the result won't drift just because you change how you order things or add some junk rows at the end.

Meng: Okay, so from a deployment angle, this means we can be more confident that if our data pipeline adds some padding to make batches uniform, the core pooling logic won't suddenly become wildly inaccurate just because of that extra padding <ref:2610.17132#pg8>.

Jane: So what’s the next thing we need to know about this paper? What are they still saying might be missing or what’s the current limitation?

Paper summary: Tom: They do mention a specific technical detail they couldn't fully certify yet. They note that a certified version of the kernel is currently needing a finite formulation of the masked maximum to be fully certified in this framework <ref:2610.17132#pg9>.

Lalam: And their final check confirms that even with floating-point sums, like when shuffling rows or appending rows with NaN, the float32 output remains the same because reordering a floating-point sum can change its last bits <ref:2610.17132#pg8>.

Jane: So, to summarize what we’ve covered about this paper, this work is a formal proof in Lean four that proves masked attention pooling is invariant to padding and permutation by showing size-generic theorems hold for the definition of the guarded normalizer <ref:2610.17132#pg5>.

Tom: And it’s connected directly to a JAX kernel implementation, verifying those mathematical invariants on specific shapes like three two and five two <ref:2610.17132#pg6>. It’s about showing the math matches the code for these critical properties.

Lu: The implication is that we have a solid foundation for using this pooling mechanism in weakly supervised computational pathology where you treat whole-slide images as bags of tile embeddings <ref:2610.17132#pg5>.

Meng: For me, the implication is confidence in the architecture. If we can mathematically prove that shuffling inputs doesn't change the output, it makes designing more complex MIL architectures much less risky during iteration <ref:2610.17132#pg9>.

Jane: So, what’s left when we look at this paper? What’s the final thought on its impact?

Tom: The development focuses only on proving structural properties of the aggregation layer; they aren't claiming accuracy or clinical validity for any specific application <ref:2610.17132#pg9>. It's a deep dive into how the attention mechanism behaves structurally <ref:2610.17132#pg6>.

Lalam: So, the main point is that if you build your pooling layer this way, you get these strong guarantees about padding and reordering resilience over arbitrary finite index types and feature dimensions <ref:2610.17132#pg5>.

Jane: That’s where we’ll leave it for now. We've talked about how this formal proof in Lean four connects the abstract math to a concrete JAX kernel, and what that means for the stability of attention-based MIL systems <ref:2610.17132#pg6>.

Tom: Right. We’ll wrap up here with that overview of "Masked Attention Pooling Is Invariant to Padding and Permutation: A Formal Proof in Lean" <ref:2610.17132#pg5>.

Conclusion: Tom: So, we've seen how they proved that masked attention pooling is stable even when you pad rows or shuffle them around using Lean four <ref:2610.17132#pg5>.

Jane: Right, so what does that actually mean for us in the real world? What’s the title of this paper, and who wrote it?

Tom: The title is "Masked Attention Pooling Is Invariant to Padding and Permutation: A Formal Proof in Lean." It was written by some researchers working with Lean four and mathlib <ref:2610.17132#pg5>.

Lu: They were focused on proving size-generic theorems regarding invariance under any bijection of the row index type, and invariance under appending any finite set of masked rows <ref:2610.17132#pg5>.

Meng: So if this holds up mathematically, what does that mean for me when I’m deploying a model? Does it mean I can stop worrying about my specific padding strategy?

Jane: It means the structure of the pooling layer itself is solid. They showed that as long as you follow their rules for guarded normalizers, the output won't change just because you added some extra rows to make your batch uniform <ref:2610.17132#pg3>.

Tom: Exactly. It means we can be more confident in how we handle data augmentation or batching without worrying about the pooling mechanism suddenly breaking due to that change <ref:2610.17132#pg8>.

Lu: The connection they made between the Lean proof and the JAX kernel implementation at shapes like five two shows these theorems hold for concrete tensor sizes <ref:2610.17132#pg6>.

Meng: And that's important because I need to know if this holds up when we run it on our actual hardware and with real traffic <ref:2610.17132#pg9>.

Jane: It does, for the shapes they tested. The implication is that we have a very strong mathematical guarantee about the aggregation step itself <ref:2610.17132#pg9>.

Tom: What's still left to look at is that they flagged one technical detail where they couldn't fully certify everything yet <ref:2610.17132#pg9>.

Lu: They mentioned needing a finite formulation of the masked maximum to be fully certified in their current setup <ref:2610.17132#pg9>.

Meng: So it’s not totally perfect yet, but it's a huge step toward knowing how robust this part of the architecture is when we scale up <ref:2610.17132#pg9>.

Jane: It’s a big step indeed. And that structural stability is what gives us confidence in building larger, more flexible models for things like whole-slide imaging <ref:2610.17132#pg5>.

Tom: Yeah, it tells us the math behind the attention mechanism is quite resilient to these kinds of data handling quirks <ref:2610.17132#pg6>.

Lu: So, next up, we’re going to look at what this structural stability means when we apply it to instance-local scorers in the paper.

Yusuf Khaled Sadek Enes Yilmaz

The Ohio State University

cs.LG

Submitted: 2026-10-12

Updated: 2026-10-12

Comments: 9 pages, 1 figure, 2 tables. Code: https://github.com/y2x0/masked-attention-lean

Code: https://github.com/y2x0/masked-attention-lean

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 89/100

The gist: The gist The paper defines and formally proves that masked attention pooling is invariant to padding and permutation using Lean 4, connecting these mathematical invariances to a JAX kernel

Key concepts

Masked Attention Pooling: This technique aggregates variable-sized instance embeddings (like tiles from an image) into a single vector. It involves calculating scores for every instance, normalizing them via softmax, and taking a weighted average of the embeddings. It is used in weakly supervised learning.
Padding Invariance: This property states that appending any number of masked rows to the input bag does not alter the resulting pooled output vector. The proof shows this because padded rows are treated as having zero mass and zero feature contribution during the aggregation process.
Permutation Invariance: This property asserts that reordering the instance embeddings, their corresponding scores, and their masks together does not change the final aggregated result. The formal proof demonstrates this by showing that reordering a set of items results in a sum over the original set plus a sum over zero-mass padded rows.
Lean 4 Formal Proof: The mathematical rigor is established using Lean 4, a formal proof assistant. This allows the authors to rigorously verify structural properties like invariance under arbitrary bijections (permutations) and appending arbitrary finite sets of masked rows, ensuring the pooling mechanism is mathematically sound.

Summary

The gist The paper defines and formally proves that masked attention pooling is invariant to padding and permutation using Lean 4, connecting these mathematical invariances to a JAX kernel implementation.

How it works

Attention-based multiple instance learning (MIL) pools a variable-size bag of instance embeddings into one vector by computing a scalar score for every instance, normalizing the scores with a softmax over the bag, and returning the weighted average of the instance embeddings <ref:2610.17132#pg2>. This construction is fundamental to weakly supervised computational pathology where a whole-slide image is treated as a bag of tile embeddings <ref:2610.17132#pg5>.

The paper addresses two key properties expected in batched implementations:

Padding invariance, which states that Appending masked rows, with whatever contents, does not change the pooled output <ref:2610.17132#pg3>.
Permutation invariance, which asserts that Reordering the rows, with each row’s features, score, and mask moving together, does not change the output <ref:2610.17132#pg5>.

Invariance Proofs in Lean

The formal proof is conducted in Lean 4 with mathlib to establish size-generic theorems regarding these invariances <ref:2610.17132#pg5>. Key theorems include:

**"Size-generic Lean 4 theorems: invariance under any bijection of the row index type, invariance under appending any finite set of masked rows with arbitrary real features and logits, irrelevance of masked values, and a zero output for all-masked bags" <ref:2610.17132#pg5>. **

Theorem 2 proves permutation invariance by showing that under a bijection e, the sum over the padded index set I ⊕ P splits into a sum over I plus a sum over P, where every padded row is masked and contributes zero mass <ref:2610.17132#pg5>.

Theorem 3 establishes padding invariance by showing that A sum over I ⊕ P splits into a sum over I plus a sum over P. Every padded row is masked, so its mass and masked feature are 0 regardless of p and t <ref:2610.17132#pg5>.

Connection to JAX Implementation

The mathematical theorems are connected to the code by tracing a JAX kernel, specifically masked pool, at shapes [3, 2] and [5, 2] <ref:2610.17132#pg6>. This tracing is facilitated by jax-lean’s certificate pipeline which translates the Jaxpr into Lean syntax <ref:2610.17132#pg6>. The paper verifies that the certified theorems hold for two concrete shapes, and that the JAX kernel meets these specifications <ref:2610.17132#pg6>.

Key Findings and Limitations

The formal definition of masked attention pooling uses a guarded normalizer D to fix the value on empty bags and avoid division by zero <ref:2610.17132#pg5>. Proposition 5 states that If no row is valid, or I is empty, the numerator is the zero vector and the output is 0 <ref:2610.17132#pg6>. Furthermore, it proves that all-masked bags maps to zero <ref:2610.17132#pg6>. The results concern ideal real arithmetic, meaning they are exact identities over R and do not model IEEE-754 rounding or XLA’s compiled code <ref:2610.17132#pg8>. The certified version of the kernel is noted as needing a finite formulation of the masked maximum to be fully certified in this current framework <ref:2610.17132#pg9>.

Corollary and Scope

A corollary covers instance-local scorers, stating that for any function f and permutation e, Poolx ◦ e, f ◦ x ◦ e, m ◦ e = Pool(x, f ◦ x, m) <ref:2610.17132#pg6>. The development focuses on proving structural properties of the aggregation layer only and does not make claims about accuracy or clinical validity <ref:2610.17132#pg9>. The paper confirms that the theorems hold for arbitrary finite index types and feature dimensions <ref:2610.17132#pg5>.

The development was built with Lean 4.33.0, the mathlib release pinned by jax-lean at revision 570aed6, and JAX/jaxlib 0.8.0 on a CPU device. The theorems concern real arithmetic, so we also test execution separately. Six JAX tests compare the kernel and a numerically stabilized variant against an independent NumPy formula. As a worked example, take X = ((1 2) (3 4) (2-1)), a = (0.2-0.1!), s = Xa = (0.5). With all three rows valid, the pooled output is (2.057208, 1.353158) in both float64 and float32. Appending two rows of NaN or shuffling the rows gives the same float32 output because reordering a floating-point sum can change its last bits. The paper confirms that masked attention pooling is padding- and permutation-invariant if masked rows contribute zero mass to the normalizer as well as the numerator, and the normalizer is guarded for empty bags.

The code availability is at https://github.com/y2x0/masked-attention-lean, commit 97978b6. The full check runs with./scripts/check.sh. Claude (Anthropic, Claude Opus 5.5) helped write and revise this manuscript from the project’s source, documentation, and validation record. The authors reviewed the text and take full responsibility for its content.

References

[1] G. Campanella, M. G. Hanna, L. Geneslaw, A. Miraflor, V. Werneck Krauss Silva, K. J. Busam, E Brogi, V E Reuter D S Klimstra and T J Fuchs.

[2] T G Dietterich R H Lathrop and T Lozano-Pérez Solving the multiple instance problem with axis-parallel rectangles Artificial Intelligence 89(1–2):31–71, 1997.

[3] D Goldberg What every computer scientist should know about floating-point arithmetic ACM Computing Surveys 23(1):5–48, 1991.

[4] M Ilse J M Tomczak and M Welling Attention-based deep multiple instance learning In Proceedings of the 35th International Conference on Machine Learning volume 80 of PMLR, pages 2127–2136, 2018.

[5] J Bradbury R Frostig P Hawkins M J Johnson C Leary D Maclaurin G Necula A Paszke J VanderPlas S Wanderman-Milne and Q Zhang JAX composable transformations of Python+NumPy programs 2018 https://github.com/jax-ml/jax.

[6] S Rush jax-lean provably correct tensor puzzles Software https://github.com/srush/jax-lean, revision 570aed6, 2026.

[7] L de Moura and S Ullrich The Lean 4 theorem prover and programming language In Automated Deduction CADE 28 volume 12699 of Lecture Notes in Computer Science, pages 625–635 Springer, 2021 doi:10.1007/978-3-030-79876-5 37 <ref:2610.

Improvements for AI systems

textbf Masked Attention Pooling Invariance for Robust MIL Training: Developed systems that pool variable-size bags of instance embeddings into a single vector while maintaining output stability regardless of padding length or instance ordering. This allows for training on large, heterogeneous datasets without needing manual padding adjustments in the batch loop, as the pooled output is expected not to depend on how much padding was added or on the order in which instances are stored.
textbf Handling Arbitrary Bag Structures: AI models can now process bags where rows are arbitrarily reindexed or appended with any finite number of masked rows without changing the aggregated result. This generalizes standard Attention-based MIL pooling to complex, dynamically shaped inputs, ensuring invariance under any bijective reindexing of rows and under appending any finite number of masked rows with arbitrary real features and logits.
textbf Robustness Against Feature Masking Errors: Systems can reliably ignore padded or irrelevant instances by utilizing the formal proof that values at masked positions are irrelevant, even when using common implementation idioms like zeroing features, which is explicitly shown to fail in real arithmetic. This prevents errors where a padded row with logit 0 still contributes e 0 = 1 to the softmax normalizer.
textbf Guaranteed Zero Output for Empty Bags: The system provides a safe fallback for empty inputs, ensuring that an all-masked bag maps to zero, and critically, it uses a guarded normalizer which fixes the value on empty bags and avoids 0/0. This eliminates catastrophic failures associated with empty bags in MIL applications.
textbf Generalization to Instance-Local Scorers: The pooling mechanism can be applied universally across different scoring functions, including complex ones like MLPs or gated attention scorers, because The corollary covering any shared, instance-local scoring function holds for arbitrary functions like a linear layer, an MLP, or a gated attention scorer.

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-infinity 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.