skillfed
RESEARCH

Checking the attention mask is not enough to guarantee causal correctness

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

Prefix invariance is a foundational guarantee for any autoregressive model: what the model computes at position t must not depend on tokens that come after t. Violating it means the model is peeking at the future during inference in ways that will not be available in real deployment. This paper formalizes that guarantee and then makes an uncomfortable argument: checking the attention mask, which is what nearly everyone does, is not enough.

The core claim is that causality is a graph-level property, not a mask-level one. A correct attention mask can coexist with a broken model if the leak travels through a different path — a scan operation, an aggregation, or a normalization layer that implicitly mixes positions. State-space and hybrid architectures introduce exactly these paths, and they are not covered by mask inspection at all.

The proposed audit is deliberately minimal: two forward passes, no training, no gradients. It produces a per-layer score that localizes where the causal violation occurs. Across 192 injected-fault trials on eight checkpoints, mask inspection detected none of the faults. The audit found all 192 of them, pinpointing the exact layer in every case. That is not a marginal improvement over the baseline — it is a categorical one.

The paper then applies the same analysis to real production models. Static and dynamic analysis of chunked-scan code in the transformers library turned up the same class of defect in Zamba2 and Nemotron-H: an inter-chunk axis error. The fix came from the reference implementation, suggesting a porting error rather than a design flaw — but the bug was present in shipped checkpoints and invisible to mask inspection.

What makes this practically significant is the cost. The audit fits on one page and runs in seconds. There is no reason not to run it before deploying any hybrid model. The field has been treating mask correctness as a proxy for causal correctness, and that proxy fails precisely for the architectures — SSMs, chunked-scan hybrids — that have attracted the most recent deployment interest.

Mask inspection is a broken proxy for causal correctness in hybrid models — this two-pass audit finds what it misses, including real bugs in shipped checkpoints.

Sources & links