Proving inference · Operator atlas · Rotary positional embedding

Rotary positional embedding

Gemma 3, Llama, most modern LLMs

$$Y_i = X_i \circ W_{1,i} \;+\; X_i' \circ W_{2,i} \\[6pt] X_i' = \big[-X_i[2],\, X_i[1],\, -X_i[4],\, X_i[3],\, \dots\big] \quad\text{(pairwise swap and negate)} \\[10pt] \text{trick:}\quad \tilde f_X(x,z) = (1-x)\tilde f_X(0,z) + x\,\tilde f_X(1,z) \\[2pt] \tilde f_{X'}(x,z) = (1-x)\big(-\tilde f_X(1,z)\big) + x\,\tilde f_X(0,z)$$

The algorithm

Rotate each pair of query/key dimensions by an angle proportional to position. Two Hadamard products against learned cos/sin matrices, with a swap-and-negate shuffle between them.

Why it resists a proof

The shuffled tensor $X'$ is a derived witness. Naively the prover must commit to it and open it, doubling the cost of a cheap operator — so the trick is to avoid ever materialising $X'$.

Cryptographic toolbox

  • Hadamard-product sumcheck
  • permuted_eq_poly: fold the swap-and-negate into the eq weights
  • Two layouts: Adjacent (pairwise) and RotateHalf (half-rows)

How each scheme proves it

DeepProve

Never commit X′ — but the mechanism is not the paper's 'two evaluations of f_X'. The code reuses a single input MLE and folds the swap-and-negate into the eq weights (permuted_eq_poly), emitting one input claim, not two.

  • Correction: our old page said the verifier reconstructs $X'$ from two evaluations of $f_X$. The code instead builds one input_mle and uses it in both sumcheck terms; the permutation lives in the eq weights via permuted_eq_poly. The sumcheck emits exactly three evals (input, cos, sin) and a single input claim.
  • The single virtual polynomial is $\text{input}\cdot\text{lt}\cdot(\text{eq}\cdot\cos + \text{permuted\_eq}\cdot\sin)$ — the swap-and-negate never touches the input tensor, only the equality polynomial.
  • There are two layouts, not one: RopeLayout::Adjacent (swap-and-negate adjacent pairs, as our old page assumed) and RopeLayout::RotateHalf (negate/rotate half-rows), which permutes differently and drops the top row-bit instead of the lowest.
  • Prover $O(sd)$, proof $O(\log sd)$. Costs 0.06% of PIOP time in GPT-2, 0.83% in Gemma 3.
zkml/src/layers/transformer/positional/rope.rs · 238-249
    fn permuted_eq_poly<F: PrimeField>(&self, eq_evals: &[F], row_size: usize) -> Vec<F> {
        match self {
            RopeLayout::Adjacent => eq_evals
                .chunks(2)
                .flat_map(|chunk| vec![chunk[1], F::ZERO - chunk[0]])
                .collect_vec(),
            RopeLayout::RotateHalf => {
                let half_len = row_size / 2;
                eq_evals
                    .chunks(row_size)
                    .flat_map(|chunk| {
                        let (first_half, second_half) = chunk.split_at(half_len);
                        second_half.iter().copied()
                            .chain(first_half.iter().map(|e| F::ZERO - *e))
                    })
                    .collect_vec()
            }
        }
    }
permuted_eq_poly folds swap-and-negate into the eq weights instead of building X′ — the actual mechanism, not 'two evaluations of f_X'.
Jolt Atlas

The trig is proven by neural teleportation: range-reduce the angle modulo τ=4π, look the reduced value up in a tiny 2¹² table, range-check the remainder. No RoPE LLM ships in the repo, so this is proven-out by unit tests only.

  • A three-stage protocol: a degree-2 division sumcheck enforcing $\mathrm{input} = \tau\!\cdot\!q + r$ with $\tau = \texttt{FOUR\_PI\_APPROX}$ ($4\pi$ scaled by 256 ≈ 3217), a one-hot lookup of the reduced value into a dense $2^{12}$ table, and a remainder range check $r \in [0,\tau)$.
  • The only algebraic tie between quotient, remainder and input is a single constraint eq · (τ·q + r − input) — everything else about $q, r$ is prover advice.
  • Scope: the repo ships only GPT-2 / nanoGPT (absolute positional), so sin/cos exist as operator support but no rotary-embedding model exercises them end-to-end.
Audit surface

Because $q, r$ are constrained only jointly by τ·q + r − input = 0, uniqueness rests entirely on the remainder range check r < τ. The remainder-range assertion at cos.rs:329 is a plain assert! that always runs in the prover — confirm the in-circuit Shout range check (not the assert) is what the verifier actually relies on.

← previousAbsolute positional encoding