Multi-head attention

Proving inference · Operator atlas · Multi-head attention

and GQA for Gemma 3

$$Y_i = \mathrm{Softmax}\!\left(\frac{Q_i K_i^\top}{\sqrt{d/h}}\right) V_i,\quad i\in[h], \qquad Y = [Y_1 \,\|\, \cdots \,\|\, Y_h] \\[10pt] \text{if } d/h \text{ is not a power of two: pad to } \tilde d_k = \text{next\_pow2}(d/h) \\[4pt] \text{and prove the padding: } \hat Y = Y\cdot P_1,\qquad P_1\in\{0,1\}^{d\times \tilde d_k h}$$

The algorithm

Run $h$ independent attention heads on column-slices of $Q,K,V$, then concatenate. Gemma 3 uses grouped-query attention, sharing $K/V$ across query groups.

Why it resists a proof

Two problems. First, the verifier holds an oracle for the concatenated $Y$, not each $Y_i$, so per-head evaluation claims must be recombined — which only works cleanly when the head dimension is a power of two, and it usually isn't. Second, every sub-operator must be batched or you pay $h$ times over.

Cryptographic toolbox

Batched everything: matmul, softmax, lookups
Random linear combination of per-head claims via β(i, r_h)
Sparse 0/1 padding matrices, MLE-evaluable in Õ(d)

How each scheme proves it

zkLLM

zkAttn runs per head; heads are concatenated and projected by another matmul sumcheck. Multi-head is composition, not a new protocol.

Each head is an independent zkAttn instance (scores → digit-decomposition softmax → ·V). The concat and output projection $W^O$ are proven as one more matmul sumcheck.
Isolated-zkAttn benchmarks show proof size and verify time sublinear in the layer count; sequence length is the dominant cost.

DeepProve

A 13-step composition that runs the whole block backwards, batching across heads, with an explicit power-of-two padding argument.

Power-of-two case. Recombine per-head claims as $y_Y = \sum_i \beta(i, r_h)\,y_{Y_i}$ for a fresh random $r_h\in\mathbb{F}^{\log h}$. Same trick for $Q, K^\top, V$.
General case. Pad each head to $\tilde d_k$ with zero columns, then prove the padding as a matmul $\hat Y = Y P_1$ against a sparse permutation matrix with exactly $d$ non-zeros. The verifier evaluates $\tilde f_{P_1}$ locally in $\tilde{O}(d)$.
Order, output → input: recombine $Y$ → batch matmul $(\hat Y_i = T_{s,i}\hat V_i)$ → batch softmax → masking PIOP → batch matmul $(T_p = \hat Q \hat K^\top)$ → recombine $Q,K,V$ → batch matmul $(Q = XW_Q)$.
Every step hands the next an evaluation claim rather than a committed tensor. Intermediate activations are never committed.
Gemma 3's grouped-query attention drops out of the same machinery: it is MHA where $K/V$ heads are shared, so the batched instance count changes but no protocol does.

Jolt Atlas

Composed from Einsum contractions and the 4-stage softmax, with no head-recombination or padding argument — and the code shows why: batching rides inside the one contraction sumcheck, so the problem DeepProve solves never arises.

Attention is named as one of three target workloads (classification, embedding generation, transformer attention), and Einsum covers the batched contractions.
The composition is Einsum scores → causal mask (a baked sentinel) → the 4-stage softmax → Einsum context. Batch matmul over heads is one EqSchedule inside a single contraction sumcheck, not a per-head loop.
So the power-of-two head-dimension problem, per-head claim recombination and explicit padding matrices that DeepProve needs genuinely do not appear — the batched einsum absorbs them. The paper doesn't spell this out; the code does.