December 6, 2022

The GKR Protocol

It has been a long time since the last post but it is finally time to play around with the last major protocol from Chapter 4 in The Book: the GKR protocol. It will involve running a protocol inside another protocol, let's go.

The GKR protocol is an interactive protocol for Arithmetic Circuit evaluation. The Arithmetic Circuit would be a circuit that has an input layer, an output layer and layers of gates in between. Gates may be addition or multiplication ones and each one of them takes exactly two inputs from the previous layer and outputs exactly one value to the next layer.

The prover $P$ starts the protocol with a claim about the output of the circuit given some input. $V$ wants to verify this claim without performing the actual circuit evaluation.

One other assumption is that the number of gates at each layer is a power of $2$.

Implementing the primitive for the Arithmetic Circuit

The concept of an Arithmetic Circuit is not a very complicated one it would make sense to start the implementation with it.

So the type of the gates would be straightforward:

pub enum GateType {
    Add,
    Mul,
}

Each gate needs to specify its type and what are the inputs from the previous layer (in our case just two indices into a slice):

pub struct Gate {
    ttype: GateType,
    inputs: [usize; 2],
}

And a layer of a circuit would just be a row of gates:

pub struct CircuitLayer {
    layer: Vec<Gate>,
}

Great, now we can describe the circuit which would hold the info about the layers and a number of inputs.

pub struct Circuit {
    layers: Vec<CircuitLayer>,
    num_inputs: usize,
}

So what info about evaluation this circuit would we need for the protocol? Not only the circuit outputs will be needed but all the values of all gates outputs at each layer. So the evaluation type may be described as such:

pub struct CircuitEvaluation<F> {
    /// Evaluations on per-layer basis.
    pub layers: Vec<Vec<F>>,
}

The implementation of circuit evaluation is trivial:

impl Circuit {
    pub fn evaluate<F>(&self, input: &[F]) -> CircuitEvaluation<F>
    where
        F: Add<Output = F> + Mul<Output = F> + Copy,
    {
        let mut layers = vec![];
        let mut current_input = input;

        layers.push(input.to_vec());

        for layer in self.layers.iter().rev() {
            let temp_layer: Vec<_> = layer
                .layer
                .iter()
                .map(|e| match e.ttype {
                    GateType::Add => current_input[e.inputs[0]] + current_input[e.inputs[1]],
                    GateType::Mul => current_input[e.inputs[0]] * current_input[e.inputs[1]],
                })
                .collect();

            layers.push(temp_layer);
            current_input = &layers[layers.len() - 1];
        }

        layers.reverse();
        CircuitEvaluation { layers }
    }
}

Since the circuit itself only describes the gates structure and is agnostic towards the types of the values it can evaluate it makes sense implementing the evaluation being generic over any type that can be Added and Multiplied.

The other two functions needed for the protocol are $\text{add}_i$ and $\text{mul}_i$ that constitute the wiring predicate of the layer $i$. It would be best to look up the definitions of these functions in The Book, but the intuition for them is this:

Each of these functions takes as input three labels:

One label of the gate in the layer $i$.
Two labels of the gates in the layer $i+1$.

Each of these functions outputs 1 in case the said two gates from $i+1$ are the inputs of the gate at level $i$ and if that is an addition or multiplication gate respectively. The protocol is going to use $\widetilde{add}_i$ and $\widetilde{mul}_i$ multilinear extensions of these functions with the first $n$-something variables fixed.

The types of methods for computing these multilinear extensions would look like this, I am omitting the code here for ergonomics:

impl Circuit {
    pub fn add_i_ext<F: Field>(&self, r_i: &[F], i: usize) ->
        DenseMultilinearExtension<F> {
        // ...
    }

    pub fn mul_i_ext<F: Field>(&self, r_i: &[F], i: usize) ->
        DenseMultilinearExtension<F> {
        // ...
    }

Other than that the Circuit would have some getter methods but that is too trivial to be included here.

The implementations of the above types may be found in circuit.rs

Building the Sum-Check polynomial

Another function that is used in the protocol is $W_i : \lbrace 0, 1 \rbrace \rightarrow \mathbb{F}$ that takes a binary gate label and outputs the gate's value at layer $i$. Again, in the protocol the MLE $\tilde{W}_i$ of $W_i$ is going to be used.

At each round of the GKR protocol an instance of Sum-Check protocol is run for the following polynomial:

$$ f^{(i)}_{r_i}(b, c) := \widetilde{add}_i(r_i, b, c) (\tilde{W}_{i+1}(b) + \tilde{W}_{i+1}(c)) + \widetilde{mul}_i(r_i, b, c) (\tilde{W}_{i+1}(b) \cdot \tilde{W}_{i+1}(c)) $$

For some fixed set of first variables $r_i$.

The definition of the above polynomial would be something like this:

pub struct W<F: Field> {
    add_i: DenseMultilinearExtension<F>,
    mul_i: DenseMultilinearExtension<F>,
    w_b: DenseMultilinearExtension<F>,
    w_c: DenseMultilinearExtension<F>,
}

The implementation of necessary trait SumCheckPolynomial for this poly is purely mechanical and the logic behind it can be found in the previous posts. The source code itself can be found at round_polynomial.rs

Restricting a polynomial to a line

One of more interesting and non-trivial building blocks of this protocol is the reduction of two evaluations of $\tilde{W}_i$ at two points $b$ and $c$ to a single evaluation at some point. This trick was described earlier in the book at section 4.5.2 and was used to shave off a constant from the verifier running time.

Here it is used since the verifier does not know $\tilde{W}_i$.

In essence suppose we have two points $b, c \in \mathbb{F}^{\log n}$. We want to create such a line $l : \mathbb{F} \rightarrow \mathbb{F}^{\log n}$ that $l(0) = b$ and $l(1) = c$. Then a multivariate polynomial, say, $\tilde{W}_i(b)$ may be restricted to line $l$.

In general case $l$ can be written as $n$ linear equations:

$$ l(x) = (k_1 x +l_1, k_2 x + l_2,\cdots,k_n x + l_n) $$

Lets mark the $i$-th component on the right side

$$ l_i(x) = (k_i x + l_i) $$

Also it is given that $l_i(0) = b_i$ and $l_i(1) = c_i$ which gives us a system of linear equations:

$$ \begin{cases} k_i \cdot 0 + l_i = b_i \\ k_i \cdot 1 + l_i = c_i \end{cases} $$

Simplifying first then substituting result into second:

$$ \begin{cases} l_i = b_i \\ k_i = c_i - b_i \end{cases} $$

And as such $l_i(x) = (c_i - b_i)\cdot x + b_i$.

For example

$b = (2, 4), c = (3, 2)$

$l(t) = ((3-2)\cdot t + 2, (2-4)\cdot t + 4)$

$l(t) = (t + 2, 4 -2\cdot t)$. $\square$

And then for polynomial $\tilde{W}(x_1,x_2) = 3x_1x_2 + 2x_2$ the restriction of $\tilde{W}$ to $l$ is $3(t + 2)(4 - 2t) + 2(4 - 2t) = -6t^2 - 4t + 32$.

Now for the first step of creating an above line $l$ the implementation is as follows:

fn line<F: Field>(b: &[F], c: &[F]) -> Vec<univariate::SparsePolynomial<F>> {
    iter::zip(b, c)
        .map(|(b, c)| {
            univariate::SparsePolynomial::from_coefficients_slice(&[(0, *b), (1, *c - b)])
        })
        .collect()
}

The second part with the restriction of the polynomial $\tilde{W}$ to this line is quite harder since it involves multiplying univariate polynomials with each other and the logistics of this in arkworks is not ideal so the code has to go back and forth between different types of univariate polynomials:

fn restrict_poly<F: Field, M: MultilinearExtension<F>>(
    b: &[F],
    c: &[F],
    mle: &M,
) -> univariate::SparsePolynomial<F> {
    let k: Vec<_> = iter::zip(b, c).map(|(b, c)| *c - b).collect();

    let evaluations = mle.to_evaluations();
    let num_vars = mle.num_vars();

    let mut res = univariate::SparsePolynomial::zero();

    for (i, evaluation) in evaluations.iter().enumerate() {
        let mut p = univariate::SparsePolynomial::from_coefficients_vec(vec![(0, *evaluation)]);
        for bit in 0..num_vars {
            let mut b =
                univariate::SparsePolynomial::from_coefficients_vec(vec![(0, b[bit]), (1, k[bit])]);

            if i & (1 << bit) == 0 {
                b = (&univariate::DensePolynomial::from_coefficients_vec(vec![F::one()]) - &b)
                    .into();
            }

            p = p.mul(&b);
        }

        res += &p;
    }

    res
}

The idea behind this code is that we first need to go from MLE in evaluation form to a polynomial in coefficients form via Lagrange interpolation. Then the resulting polynomial must be restricted to line $l$. However, I've found that the easiest way to do the above two steps with arkworks is to do them at the same time and sort of interpolate directly into the restricted univariate polynomial.

The rest of the protocol

The remainder of the protocol implementation is purely mechanical and involves wiring together the Sum-Check protocols that are run at each layer. The mechanical pages are no fun to describe plus I am not quite happy how they turned out. Maybe I will revisit the messaging implementation in the future and make it the topic for a followup post. The code of the protocol can be found in lib.rs. Thanks for reading and stay tuned!

Reference list.

Original paper by Shafi Goldwasser, Yael Kalai and Guy Rothblum

YouTube: GKR based Zero-Knowledge Proofs - Yael Kalai, Microsoft Research

Top