Proving inference · Operator atlas · Matrix multiplication

Matrix multiplication

QKV projections, FFN, logits

$$D = a\,(A\cdot B) + b\,C, \qquad A\in\mathbb{F}^{m\times l},\; B\in\mathbb{F}^{l\times n} \\[10pt] \text{sumcheck instance:}\qquad y_D - b\,y_C \;=\; a \!\!\sum_{x\in\{0,1\}^{\log l}}\!\! \tilde{f}_A(r_1,x)\cdot \tilde{f}_B(x,r_2) \\[6pt] \text{prover } O(ml+ln+mn), \qquad \text{proof } O(\log l)$$

The algorithm

The workhorse: $Q = XW_Q$, the FFN's two dense layers, the final unembedding. Quantized, $x_4[i,k]=z_4+M\sum_j(x_1[i,j]-z_1)(x_2[j,k]-z_2)+B(x_3[k]-z_3)$ with $M=s_1s_2/s_4$.

Why it resists a proof

It does not resist. Sumcheck proves it with a linear-time prover, a logarithmic proof, and no commitment to the product matrix. Only the inputs need oracles; their claims pass down the chain. The danger is treating it as a generic circuit, which costs large constant factors.

Cryptographic toolbox

  • Thaler's matmul sumcheck
  • Batched sumcheck across heads / layers
  • Virtual polynomials: never commit intermediates
  • Verifying tensor relations directly at the multilinear level (Schwartz–Zippel)

How each scheme proves it

Thaler's matmul sumcheck, made ~10× faster by a grouping algorithm that turns the prover's bottleneck from field multiplications into field additions.

  • Thaler's protocol is bottlenecked by building each matrix's bookkeeping table (~4nm field multiplications). zkGPT's grouping algorithm halves that, then exploits two properties of quantized LLM matrices — the padded zeros sumcheck requires, and the small value range from 16-bit quantization — to shift the bottleneck to n·m field additions.
  • Additions are several times cheaper than multiplications, so the net is ~10× on a protocol that was already theoretically optimal. This is the linear-layer half of "prove GPT-2 in <25 s."
  • Circuit squeeze adds a companion protocol that checks multiple matmuls in one sumcheck, so the matmul layers (proven by Thaler's protocol, not GKR) can also be parallelised.

The standard sumcheck for tensor multiplication, GPU-parallelised. The arithmetic backbone under Q·Kᵀ, ·V, and the projections.

  • $C=AB$ is one sumcheck over the multilinear extensions, reduced to evaluation claims on $\tilde A, \tilde B, \tilde C$ and closed by Hyrax openings.
  • Tensor structure is preserved, so the sumcheck parallelises across dimensions. That is the design constraint that makes GPU proving of a 13B model feasible.

Generalised, batched matmul-with-bias sumcheck; the primitive everything else reduces to.

  • A single batched general matmul PIOP handles $h$ instances $a_iA_iB_i + b_iC_i = D_i$ with one sumcheck. Used for all attention heads, and for the padding and permutation matrices elsewhere.
  • Round-by-round soundness error $h(1+2\log \max(n_i)\max(m_i))/|\mathbb{F}|$, composed with the sumcheck's own — this is what licenses Fiat–Shamir via the BCS transform.
  • Padding and re-arrangement (needed when $d/h$ isn't a power of two) are themselves proven as matmuls against sparse 0/1 matrices whose MLEs the verifier evaluates locally in $\tilde{O}(d)$.
  • Measured at only 14.9% of PIOP time. The classical ZKML intuition that matmul dominates is wrong at LLM scale.
Einsum / MatMul: 14.91% of PIOP time (GPT-2, seq 512).

ONNX Einsum, verified directly at the polynomial level — no decomposition into scalar operations, no R1CS.

Instead of naively decomposing tensor operations into scalar computations and verifying each element independently, we verify tensor relations directly at the multilinear polynomial level.

Jolt Atlas · §1.2

Nearly all verification checks are performed using virtual polynomials, and the approach neither uses nor requires R1CS.

Jolt Atlas · §5.4
  • One engine, all patterns. einsum/dot.rs proves ~11 einsum patterns; each supplies three hooks via the EinsumLayout trait: fold (collapse the non-contracted output axes into eq-weights at the output randomness), an EqSchedule (where a batch eq-poly rides in the round order), and operand_points (scatter the final point back to the two operand openings). Attention Q·Kᵀ and ·V are the same code path as a feed-forward matmul.
  • After folding, a matmul is a degree-2 dot product over the contraction cube: $\sum_h L(h)R(h)$ (EqSchedule::None), or degree-3 $\sum_h L(h)R(h)\,\mathrm{EQ}(h)$ when a batch eq-poly rides along (High/Low bind the batch bits first or last) — one sumcheck covers batch matmul, no per-head loop.
  • The initial claim is not the output opening. For quantized nodes fused_input_claim reconstructs $\mathrm{acc}(r)=\mathrm{rescaled}(r)\cdot 2^S + R(r)$ from two prover advice openings (ClampAcc, RescaleRemainder) — the divide-by-$2^S$ is fused into the matmul sumcheck, never a separate Div node.
  • The outer-product pattern m,an->a1nm is a degenerate zero-round instance: num_rounds()==0, it just claims the product of two point evaluations. And the shipped pipeline still lowers ONNX to a CPU trace first, so the "no R1CS" tensor-native path is a design in transition.
onnx_proof/fused_rebase.rs · 122-135
pub fn fused_input_claim<F: JoltField>(...) -> F {
    match rebase_bits(&node.operator) {
        Some(bits) if bits > 0 => {
            let rescaled = accessor.get_advice(VirtualPoly::ClampAcc).1;
            let remainder = accessor.get_advice(VirtualPoly::RescaleRemainder).1;
            rescaled * F::from_u64(1u64 << bits) + remainder
        }
        _ => accessor.get_reduced_opening().1,
    }
}
The matmul sumcheck's initial claim is reconstructed from two prover-supplied advice openings, not the output opening — the central soundness seam.
Audit surface

Operand values L, R are opened via the accumulator but are not range-checked in the einsum module — no field-overflow guard on acc lives here; it relies on upstream quantization invariants. Uniqueness of the (rescaled, R) split rests entirely on the R ∈ [0,2^S) range check being complete: a weak one lets a prover shift value between quotient and remainder and forge the clamped output. Two concrete gaps: non-power-of-two shapes are behind #[ignore]d tests ("not fully validated yet"), and in ZK mode the remainder claim is a placeholder that "reaches the proof only through the prover-supplied baked initial claims of the following stages" (fused_rebase.rs:305) — a documented trust seam.

A generalised Freivalds check: prove $u^\top ABv = u^\top Cv$ for random $u,v$ instead of $AB = C$. The circuit collapses to a single inner product, with communication independent of the inner dimension -- but the prover still multiplies the matrices.

Directly proving matrix multiplication in zero knowledge would require O(n³) number of multiplications, which could be improved to O(n^2.8) (or even lower) based on a better algorithm. Although the prover time has to be linear to this complexity, we could reduce the circuit size for ZK proofs significantly.

Mystique · §3.3
  • $\Pi_{\mathrm{MatMul}}$: the verifier samples $u \in (\mathbb{F}_{q^k})^n$, $v \in (\mathbb{F}_{q^k})^\ell$; both parties locally fold the authenticated matrices into $[x]^\top = u^\top[A]$, $[y] = [B]v$ and $[z] = u^\top[C]v$, and prove only the single inner product $[x]^\top[y] = [z]$. Soundness error $3/q^k$.
  • With the multivariate-polynomial optimisation, communication drops to $O(k\log q)$ bits -- independent of the contracted dimension $m$. The linear layer is essentially free on the wire.
  • The saving is in circuit size, not prover work. The prover must still compute $A\cdot B$ locally to know $C$, and the paper names that local multiplication as the efficiency bottleneck. Read the headline improvement as a statement about communication.
  • The same random-linear-combination idea recurs across the whole field: it is Freivalds' algorithm, and it is what the sum-check matmul PIOP is doing too -- reduce a matrix identity to a scalar identity at a random point, and never commit the product.