Source-linked AI summary
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
TL;DR
Autoregressive models rely on prefix invariance, yet releases rarely provide evidence that implementations are causally correct, and silent leakage can improve development metrics. The paper introduces a lightweight, layer-localizing audit and finds that attention-mask inspection misses injected faults while the audit detects them and identifies defects in Zamba2 and Nemotron-H.
Problem
Released autoregressive models rarely provide evidence of causal correctness, although training, evaluation, decoding, KV-cache reuse, and speculative decoding assume prefix invariance.
Method
The paper audits prefix invariance with two forward passes, no training or gradients, and per-layer hooks that localize the faulty layer rather than only detecting leakage.
Results
192/192 injected faults were localized by the audit, while static attention-mask inspection flagged 0 of 192; dynamic testing confirmed causal defects in Zamba2 and Nemotron-H.
Takeaways & Limitations
Prefix invariance requires auditing the whole computation rather than inspecting attention masks, especially in stacks combining masked and unmasked mixers.
Takeaways & Limitations
The audit did not cover the fused CUDA-kernel implementation used when optional fused dependencies are installed, so its scope is limited to the PyTorch chunked-scan path.
Abstract
from arXiv · showhide
We formalize prefix invariance: representations at position t must not depend on future inputs. We give a lightweight audit, two forward passes, no training or gradients, that localizes exactly where causality breaks. Attention-mask inspection is incomplete: leaks can occur via scans or normalization despite correct masks. Across 192 injected-fault trials on eight checkpoints, mask inspection found none, while our audit localized all 192/192, also finding a defect in Zamba2 and Nemotron-H.
1 Introduction
Autoregressive models require outputs at position t to depend only on inputs through t, yet releases rarely provide evidence that implementations satisfy this property. Because modern stacks combine heterogeneous mixers and leakage can improve conventional metrics, mask inspection alone is not a sufficient causal audit.
- Autoregressive correctness requires the output at position t to depend only on inputs at positions ≤t.
- Contemporary model releases report standard scale and benchmark information without evidence that the released implementation is causally correct.
- Mask inspection became a de facto causality audit for decoder-only Transformers because self-attention’s causal behavior is governed by an explicit mask.
- Modern sequence architectures combine attention with recurrent, state-space, and convolutional mixers, so explicit attention masks cover only part of the computation graph.
- Leakage can lower training loss and validation perplexity and improve teacher-forced benchmark scores despite violating autoregressive inference assumptions.
- The paper frames prefix invariance as a deterministic metamorphic relation and applies it to attention-only, state-space, recurrent, and hybrid architectures.
2 Method
The audit compares two forward passes that differ only at the final token, captures every layer’s output, and reports the first layer where earlier representations differ. It avoids training, gradients, labels, and accelerator dependence while preserving exact-zero behavior as the clean criterion.
- The audit creates two length-T sequences identical except at position T −1, then compares their per-layer representations on the unchanged prefix.
- The procedure returns LEAK and the first offending layer when a per-layer difference exceeds threshold τ; otherwise it returns CLEAN.
- Per-layer hooks localize faults that logits-only checks can detect but cannot localize.The logits-only variant detected every injected fault and localized none.
- The comparison excludes the final position because it is the deliberately perturbed token, isolating the prefix that should remain invariant.
- Caching is disabled so both passes follow structurally identical full-sequence paths rather than reusing state across calls.
- Deterministic single-precision evaluation makes clean per-layer differences exactly zero and the verdict robust across τ ∈[10−6, 10−3].
- The audit requires two forward passes and no backward pass, gradient graph, optimizer state, training data, labels, or accelerator.At T = 48, captured-activation memory overhead is negligible relative to model weights.
3 Experiments
The audit exactly localized every injected fault across diverse architectures, while several existing practices either missed leaks or could not identify their layer. Clean evaluations produced exact-zero deltas, but sensitivity depended on precision, threshold, sequence length, and checkpoint.
- 192/192 injected faults were localized to the exact layer across eight public checkpoints, with no checkpoint scoring below 24/24.The trials covered eight fault patterns at three depths per model.
- Mask inspection detected 0/96 injected faults because all faults altered layer outputs while leaving causal mask attributes correct.This also fails for architectures with no attention-mask attribute.
- B3 missed 25/96 faults structurally when radius-1 leaks could not reach the shuffled suffix, while one additional miss was caused by its threshold.The structural blind spot cannot be fixed by threshold recalibration.
- Clean repeated evaluations were bit-identical with zero differing elements and zero maximum absolute deviation, making verdicts threshold-robust on the tested set.The reported threshold range was τ ∈[10^-6, 10^-3].
- Audit sensitivity had explicit operating limits: float32 lost sub-resolution leaks at ε ≤1e-11, while checkpoint-specific exact-localization floors varied and LFM2-1.2B was an outlier.The authors state that leaks weaker than roughly 1e-10 relative magnitude require float64 and that floors should be measured rather than assumed.
- The eight-checkpoint census was clean at T = 48, but this does not establish causal correctness in general because some defects emerge only beyond the tested sequence length.The census nonetheless supported applicability across masked attention, state-space scans, recurrences, and hybrids without modification.
- Three Falcon-H1 checkpoints initially appeared clean with max |∆| = 0.0 at every layer, motivating a methodological change before treating such results as conclusive.This negative result concerned the paper’s own evaluation procedure.
4 Limitations
The paper identifies important scope, coverage, and comparison limits: the census is small and non-random, some models are unauditable, and several controls do not generalize across checkpoints or execution settings.
- Two released models were found to violate causality, but they share one code lineage and therefore do not establish a base rate.Zamba2 received more extensive validation than Nemotron-H, so the evidence is not fully symmetric.
- The audit localizes faults but does not improve detection over existing logits-only checks, while a gradient-based check matches its localization at higher cost.The gradient-based comparison also depends on differentiability and does not make zero gradients equivalent to absent dependence.
- The authors’ own releases remain only partially audited, with the instruction-tuned variant having only a negative control and two internal checkpoints unaudited.The primary release passed controls with 49 clean layers and 16/16 positive-control localization.
- Roughly twenty checkpoints were selected for architectural diversity rather than randomly sampled, so the census cannot support proportions or base-rate claims.The sample spans 129M to 9B parameters and includes both a clean 9B model and an 8B leaker.
- Coverage of custom-kernel hybrids is systematically thinner because some checkpoints require GPU-only kernels, are gated, or cannot load through the standard interface.The paper notes that this gap may matter because custom-kernel implementations are plausibly more bug-prone.
- Several released-model clean verdicts lack the positive-control gate because some layer return signatures reject the required output injection.The conformant reference construction does satisfy the gate, but that does not extend automatically to every released checkpoint.
5 Conclusion
The paper presents prefix invariance as a broadly necessary property and proposes a lightweight audit that localizes violations beyond attention masks. It also limits its claim: the contribution is localization, while existing methods can match detection or localization under different costs.
- Two forward passes without training, gradients, labels, or an accelerator test prefix invariance and identify the first layer where causality breaks.The proposed certificate includes per-layer prefix deltas and a same-checkpoint positive-control result.
- Static attention-mask inspection flagged 0 of 192 injected faults, demonstrating that masking one operator is not a sound audit of the whole computation graph.For two of eight audited architectures, there was no attention-mask object to inspect.
- A static census predicted Zamba2 and Nemotron-H defects from the source, and dynamic auditing confirmed contamination at their chunk boundaries.The reported boundary onsets were 256 and 128, respectively, and the root cause was reduction over the wrong axis.
- The paper does not claim better detection than a logits-only check or better localization than a gradient-based check; its distinct contribution is layer localization with lower operational cost.Two additional small deltas were rejected after perturbation sweeps showed flat responses.
- The instrument itself required positive-control safeguards after three checkpoints produced false clean verdicts, and the census length initially hid a real defect.The authors therefore connect audit validity to liveness evidence and sequence lengths exceeding internal chunk, window, and kernel parameters.
- The proposed release norm is a causal-correctness certificate reporting per-layer deltas, threshold, precision, relative sequence length, and a same-checkpoint positive control.The certificate is proposed as release information alongside the parameter count.
A.1 Why we print the method instead of shipping a package
The authors release reproducible measurements rather than a software package, arguing that independent reimplementation is preferable to running a bundled binary.
- The method consists of five lines of arithmetic added to standard forward hooks and can be reimplemented against models with enumerable layers.The authors estimate that a competent reader can reproduce the core procedure in under an hour.
- The release contains audit logs, exact checkpoint identifiers, per-layer delta arrays, injected-fault specifications, and environment details, but no implementation license or software distribution.These materials are intended to support reproduction or contesting of reported numbers.
A.2 Environment
The experiments used CPU for the census and injection suite, with a single GPU reserved for memory capacity during larger-model audits, and primarily used float32 evaluation.
- The census and injection suite ran on an Intel Xeon Gold 6526Y CPU with 64 threads and 251 GB RAM, while larger-model audits used one GPU only for memory capacity.The test itself does not require an accelerator.
- Unless otherwise stated, audits used float32, PyTorch 2.10.0, standard model-hub loading with remote code when required, evaluation mode, and disabled caching.Fused causal convolution and Mamba SSM kernels were absent, so state-space models used reference implementations.
A.3 Injected-fault specifications
The injected-fault patterns wrap a single layer’s forward pass and transform its output tensor. Faults are placed at three specified depths, with ε = 1.0 except in the Section 3.5 sweep.
- Each injected-fault pattern wraps one layer’s forward pass and transforms its output tensor o with shape (B, T, d).The definitions are exactly reproducible.
- Faults are injected at layer 1, ⌊L/2⌋, and L −2.
- Injection strength is ε = 1.0 except in the Section 3.5 sweep.
A.4 Baseline specifications
The baselines test output sensitivity, mask declarations, shuffled-suffix effects, future-token resampling, and cache consistency. They range from static inspection to repeated forward comparisons, with released scripts and result files documenting the evaluations.
- B1 compares final logits over positions [0, T −1) for the same two inputs and flags max differences above τ.By construction, B1 provides no layer information.
- B2 statically enumerates modules and flags any causal or is-causal attribute set false without running a forward pass.
- B3 shuffles positions ≥T/2 and compares cross-entropy on targets ¡ T/2 using a 1e-4 threshold.
- Table 13 covers the primary audit harnesses and baseline implementations, while Table 14 maps released result files to census audits and failure analyses.
- B6 compares incremental cached decoding with a single full-sequence forward using a 1e-4 threshold and T+1 forward passes.
A.5 Known defect in a superseded artifact
A superseded sensitivity-harness version cast hidden states to float32 before computing deltas, truncating float64 measurements; the analysis was corrected and made auditable.
- The early harness cast hidden states to float32 before computing deltas, quantizing float64 measurements to powers of two.The superseded file remains retained and marked in Appendix B.
- Section 3.5 uses only post-correction data computed in float64.
- The cast was lossless for float32 results, which were unaffected and re-verified.
Appendix B — Raw data index
The raw-data index links the evaluation’s scripts and released JSON artifacts to the reported audits, baseline comparisons, bit-level verification, and precision-sensitivity analyses.
- All experimental paths are stored on the compute host, and every Table 3 number is traceable to a retained file.Per-record data include full per-layer delta arrays, wall time, and peak RSS.
- Table 13 indexes scripts for the census harness, baseline implementations, and precision-sensitivity analyses.
- Table 14 indexes released JSON artifacts for successful public-checkpoint audits, load failures, and internal validation runs.
- Table 15 indexes JSON artifacts for baseline comparisons and bit-level verification experiments.
- Table 16 indexes artifacts for the precision and threshold sensitivity study, including exact-localization floors across precisions and thresholds.