Proving inference · Operator atlas · Causal masking

Causal masking

the −∞ upper triangle

$$\mathrm{Mask}(T) = T \circ D_1 + D_2 \\[6pt] D_1 = \text{lower-triangular ones},\qquad D_2 = \text{upper-triangular } -\infty \\[10pt] \tilde f_{D_1}(r_1,r_2) = \sum_i \beta(i,r_1)\,\tilde f^{(i)}_{D_1}(r_2), \qquad \tilde f^{(i)}_{D_1}(r_2) = \tilde f^{(i-1)}_{D_1}(r_2) + \beta(i,r_2)$$

The algorithm

Zero out attention to future tokens. This is the operator that makes the certification trick sound: the prediction at position $t$ provably cannot see token $t{+}1$.

Why it resists a proof

Naively the verifier evaluates two $s\times s$ structured matrices' MLEs — $O(s^2)$ work, which would destroy succinct verification.

Cryptographic toolbox

  • Hadamard sumcheck with an additive correction
  • Closed-form lower-triangular MLE ⇒ O(log s) verifier (causal)
  • An extra sumcheck only for the local sliding-window mask

How each scheme proves it

A single Hadamard-style sumcheck — but the causal-mask verifier is O(log s), not O(s). The lower-triangular mask MLE is evaluated in closed form (eval_zeroifier_mle); only the local sliding-window mask needs an extra sumcheck.

  • Correction: our old page said the verifier sweeps the triangular MLE in $O(s)$ via a prefix sum. The code evaluates the lower-triangular (zeroifier) MLE and the row/column less-than polynomials in closed form — the full causal-mask verifier is $O(\log s)$.
  • The verifier negates $D_2$'s $-\infty$ contribution first, reducing to a plain Hadamard check against the triangular $D_1$. It slots into single-head attention between the score matmul and the softmax.
  • The $O(s)$ cost only appears for the local sliding-window mask (Gemma-style), which needs an extra sumcheck plus an $O(s)$ shift-matrix evaluation. The plain causal mask does not.
  • zkGPT's evaluated model has no attention mask at all.
Attention mask: 3.96% (GPT-2) / 0.97% (Gemma 3) of PIOP time.

There is no masking PIOP because the mask is a baked constant: future positions carry a large-negative quantized sentinel that flows through the softmax as ≈0. Nothing separate is proven.

  • In the traced model the causal mask is applied before quantization (att.masked_fill(bias==0, -10) in models/transformer/gen.py), so by the time the score matrix reaches the prover the upper triangle is already a constant.
  • The softmax test layout documents the exact sentinel: upper-triangular entries use -(11 << 12) = -45056 (a $-\infty$ proxy at scale 12). It is just a very small input to the exp lookup, which returns ≈0 — no triangular-MLE argument, no Hadamard correction.
  • So Jolt Atlas needs neither DeepProve's prefix-sum triangular trick nor a masking cost line: masking is absorbed into the softmax it already proves.
Audit surface

Because masking is a trace-time constant, its correctness is not proven — it is trusted as part of the committed input graph. That is sound for a fixed causal mask compiled into the model, but an auditor should confirm the mask constants in the committed graph actually match the claimed causal pattern (nothing in the proof forces the upper triangle to be the sentinel rather than a real score).