Proving inference · Operator atlas · Unembedding → logits

Unembedding → logits

h · Wᵀ over the vocabulary

$$\text{projection:}\quad \mathrm{logits} = h \cdot W_U^\top, \qquad W_U \in \mathbb{F}^{V\times d} \\[6pt] \text{then:}\quad \text{next token} = \arg\max_v \mathrm{logits}[v] \\[6pt] V \approx 50{,}000 \;(\text{GPT-2}) \qquad V \approx 262{,}000 \;(\text{Gemma 3})$$

The algorithm

Project the final hidden state onto the vocabulary, then pick the next token. The paper's cost table bundles these as "logits," but in DeepProve's code they are two distinct layers, and only the first is a matmul.

Why it resists a proof

The projection is trivial but heavy: the vocabulary dimension is enormous, so this matmul and its committed output dwarf any attention layer. For Gemma 3's 262k vocabulary it is the second-largest witness in the trace. The selection that follows is the interesting part, and it is not a matmul.

Cryptographic toolbox

  • Matmul sumcheck for the projection
  • A separate argmax layer for the selection
  • LogUp range-check + Hadamard sumcheck (the argmax)

How each scheme proves it

DeepProve

Two layers, not one. The vocab projection is an EinSum (an ordinary matmul); the layer literally named Logits is the argmax head — proven with a LogUp lookup plus a Hadamard sumcheck, not a matmul.

  • Our old page called this "just a matmul." That is true only of the projection, which is an EinSum node. The code's Logits layer is the argmax that selects the output token, and its LogitsProof carries a logup_proof (a range-check lookup) plus a hadamard_proof sumcheck.
  • This is also incidental evidence for the Gather debate: the repository does contain LogUp/GKR lookup arguments — the argmax head is one of their users.
  • For Gemma 3 the projection is 18.1% of witness-commitment time and 16.1% of PIOP time — roughly double GPT-2's share, tracking the vocabulary ratio. zkGPT's evaluated model omits the final projection entirely.
zkml/src/layers/transformer/logits.rs · 163-172
pub struct LogitsProof<F: PrimeField, PCS: CommitmentScheme> {
    logup_proof: LogUpBatchProof<F>,
    /// Commitment to the MLE of the vector of maximum values
    commitments: Vec<VerifierCommitment<PCS>>,
    /// Proof of hadamard product sum-check
    hadamard_proof: IOPProof<F>,
    /// Evaluation of the input tensor MLE and max values MLE in this order
    #[serde(with = "dp_crypto::serialization")]
    sumcheck_evals: Vec<F>,
}
The Logits layer is an argmax, not the vocab matmul: its proof is a LogUp lookup plus a Hadamard-product sumcheck.
Logits (projection): 9.19% (GPT-2) / 18.11% (Gemma 3) of witness commit.
Jolt Atlas

The vocabulary projection is just another Einsum — the same fold-and-dot engine as any matmul. The only special thing is the shape.

  • $h\cdot W_U^\top$ is one more einsum/dot.rs contraction; there is no dedicated logits protocol.
  • The cost is the vocabulary dimension: GPT-2's 50257 is padded to the next power of two (65536), so the traced output is [1, 16, 65536] — the projection and its committed output dwarf any single attention layer, exactly as on the DeepProve side.
← previousGeLU activation