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

Yusuf Khaled Sadek; Enes Yilmaz

cs.LGOct 12, 2026

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

Authors: Yusuf Khaled Sadek, Enes Yilmaz

Organizations: The Ohio State University, Columbus, OH, USA

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.

Figures & tables

Idiom	Zero pad	NaN pad	Empty	Exact in $\mathbb{R}$
Zero padded features only	$(1.3563,0.8921)$	$(1.3563,0.8921)$	$(0,0)$	no
Logits $\to-\infty$ , features unmasked	correct	NaN	NaN	n/a
Logits $\to-10^{9}$ , features unmasked	correct	NaN	$(2,3)$	no
Softmax of $-\infty$ -masked logits, features zeroed	correct	correct	NaN	n/a
Definition 1 ( masked_pool )	correct	correct	$(0,0)$	yes

Table 1: Behaviour of common masking idioms in JAX (float32). The bag is the three-instance example of Section 6 , whose correct output is

(2.0572,\,1.3532)

. “Zero pad”: two appended rows with zero features and zero logits. “NaN pad”: the same rows with NaN features. “Empty”: two rows, both masked, with features

(1,2)

and

(3,4)

. The last column asks whether the idiom is exactly padding-invariant in real arithmetic.

Panel of Figure 1 — Figure 1: Chain of results. Shaded, dashed boxes and the arrows between them are trusted: JAX tracing and the jax-lean importer that encodes the Jaxpr as Lean syntax. The vertical and lower arrows are theorems checked by Lean. The certified_* theorems compose them to state both invariances directly about the imported IR.

File	Lines	Content
Basic.lean	79	Definition 1 , Theorems 2 – 3 , Props. 4 – 5 , Cor. 6
Connection.lean	79	specification lemmas, concrete and certified invariances
Generated3.lean	200	emitted function, imported IR, certificates for $[3,2]$
Generated5.lean	200	the same for $[5,2]$

Table 2: Size of the development (lines, including comments).

Multiple Instance Learning

Multiple Instance Learning learns from labeled sets of instances when labels apply to bags, enabling prediction without dense annotations.

See all 96 papers

The Mask Is Not the Model: Auditing Prefix Invariance in Attention, State-Space, and Hybrid Sequence Models
Aug 24, 2026Taebong Kim, Youngsik Hong, Minsik Kim +3LLM Auditing Sequence Modeling
VIDRAFT AI Research · QuantumOS, Seoul, Republic of Korea
MLIPS: Multi-Layer Indexer Pool Sharing
Oct 20, 2026Samkit Jain, Sai Kishan Pampana, Shlok Agarwal +4LLM Inference Acceleration Sparse Attention
NVIDIA Corporation
Power law graph attention: exact generalization of scaled dot-product attention, empirical collapse at inference
Aug 10, 2026Burc GokdenGraph Attention Networks LLM Inference
Fromthesky Research Labs LLC, Oregon, USA

Aug 24, 2026J/K move · Enter open · S save

The Mask Is Not the Model: Auditing Prefix Invariance in Attention, State-Space, and Hybrid Sequence Models

Taebong Kim, Youngsik Hong, Minsik Kim, Sunyoung Choi+2

VIDRAFT AI Research · QuantumOS, Seoul, Republic of Korea

LLM Auditing Sequence Modeling State Space Models

Open paper AskPDF arXiv

Hybrid sequence models must satisfy prefix invariance: representations at position t must not depend on future inputs, yet this is rarely verified. We formalize prefix invariance and give a lightweight audit, two forward passes, no training or gradients, yielding a per-layer score localizing where causality breaks. Attention-mask inspection, the field's default check, is incomplete: causality is a graph-level property, and leaks can occur via scans, aggregations, or normalization despite correct masks. Across 192 injected-fault trials on eight checkpoints, mask inspection detected none, while our audit localized all 192/192 to the exact layer. Static/dynamic analysis of chunked-scan code in transformers found the same defect in Zamba2 and Nemotron-H, an inter-chunk axis error fixed via the reference implementation. The method fits on one page and runs in seconds.

Abstract

Multiple Instance Learning

Explore similar work

The Mask Is Not the Model: Auditing Prefix Invariance in Attention, State-Space, and Hybrid Sequence Models

MLIPS: Multi-Layer Indexer Pool Sharing

Power law graph attention: exact generalization of scaled dot-product attention, empirical collapse at inference