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, Jaewon Jang, Minseo Kim
cs.LG, cs.AI
2026-08-24
Two forwards audit prefix invariance: change the last token, see which layer first moves. Masks caught 0/192 injected faults; it localized 192/192, including Zamba2 and Nemotron-H.
An autoregressive model has one hard constraint: the representation at position t cannot see inputs after t. Teacher forcing, KV caches, and speculative decoding all assume this. The usual check is whether the attention mask looks causal. That is no longer enough. Mamba and RWKV have no mask at all. Chunked scans, global normalization, and differential-attention aggregation can leak the future while the mask stays correct.
Leakage also makes training metrics look better. Future tokens make next-token prediction easier, so loss and teacher-forced perplexity fall together, and the break shows up only in free-running generation or the cached path. The authors hit three unrelated leaks in an internal hybrid stack. None of them moved ordinary training curves.
Prefix invariance: if two sequences share a prefix, every layer's representation on that prefix must match. The cheapest instance changes only the last token.
The test is two forward passes. Draw a random sequence of length T, copy it and replace the last token, disable the cache, hook every layer. Compare the max absolute difference on positions [0, T−1). The first layer whose delta exceeds τ (default 1e-6) is the leak origin. On a correct graph that difference is exact zero, so any τ between 1e-6 and 1e-3 gives the same verdict on their suite. A 135M model takes 0.1–0.5 s on CPU after loading.
Four design choices matter. Layer outputs, not final logits, are what send an engineer to a line of code. The last position is supposed to change, so it is excluded. Enabling the cache takes a different code path, so the test would no longer be about the graph under test. A clean verdict is admissible only if a positive-control injection fires on that same loaded checkpoint; otherwise all-zero deltas can mean the model is ignoring its input.
Faults fall into four classes and eight patterns: a one-step shift, time-axis pooling, normalization leakage, and reverse scans or last-position broadcasts. Audit length must exceed chunk, window, and kernel sizes, or those paths never run in general form. Their first census used T=48; Zamba2's chunk is 256, so the defect was invisible.
Eight public checkpoints, eight patterns, three depths: 192 injections, all localized to the injected layer. Two of those models, mamba-130m and rwkv-6-world-1b6, have no attention mask anywhere.
On 96 matched trials, static mask inspection scored 0/96. Logits-only perturbation detected 96/96 and localized none. Shuffled-suffix perplexity got 71/96 and is structurally blind to off-by-one leaks whose radius never reaches the shuffled region. Gradients of prefix outputs with respect to future positions localized 96/96, at the cost of a backward pass and a differentiable path. Cache-versus-full-sequence consistency false-alarmed on clean SmolLM2-135M with Δ=1.85×10⁻⁴.
A static census of chunked-scan code in transformers 5.7.0 flagged two of six implementations, Zamba2 and Nemotron-H, for reducing over the output-chunk axis instead of the input-chunk axis. Those three lines are byte-identical. Dynamic audits matched the prediction. Zamba2-1.2B starts contaminating at position 256, max prefix delta 1.12×10⁻². Nemotron-H-8B stays exact-zero through T=128 and jumps to 0.74 past the boundary, first at layer 17, the first Mamba whose scan crosses a chunk. Randomly initialized weights leak at the same boundary. Replacing the block with the Mamba2 reference drives the delta to exact zero. Conformant relatives such as Bamba-9B and Falcon-Mamba-7B stay clean. The leak is in the pure PyTorch path; fused CUDA kernels were not audited.
Three Falcon-H1 sizes loaded into a state that emitted bit-identical outputs for different inputs, with positive controls at 0/24. Those are reported as not audited, not as clean.
Hybrid stacks will keep mixing masked and unmasked mixers, and mask inspection covers less of the graph each time. The audit is cheap enough for CI: two forwards, no training, no GPU required for the test itself. It does not beat logits perturbation at detection. It answers which layer broke. The release recommendation is a causal-correctness certificate next to parameter counts and context length, with a live positive control and a sequence longer than the model's internal block sizes.
This is an existence proof, not a defect rate. Both leaks come from one copied axis error.
The authors rank these themselves. Detection is not better than a logits perturbation. Gradient localization matches accuracy; the case for the forward-only test is cost and applicability. The census is about twenty non-random checkpoints from 129M to 9B. Fused-kernel hybrids, a gated repo, and a missing modeltype could not be audited, which is where bugs are plausibly more likely. Some clean verdicts are ungated because layer return signatures reject the injection wrapper. LFM2-1.2B localizes all 24 injections at strength 1.0 and names a layer three steps downstream for every sweep point at strength 0.1 or below, with no mechanism given. Training mode, sharded execution, and input-dependent leaks are untested. The full alternative-explanation battery run on Zamba2 was not repeated identically on Nemotron-H.