Generalizing Sum Check protocol and counting the triangles

Welcome back. In the previous post we have taken a first look at the Sum Check protocol from The Book and implemented it for the case of multivariate::SparsePolynomial. In this post I am going to attempt to generalize the protocol to any polynomial and apply it to the Counting Triangles in Graphs problem.

In the previous post an implementation of the Sum Check protocol for the special case of multivariate::SparsePolynomial was described. But if you have been reading The Book or following the online book reading group you may already know that for the most interesting cases Sum Check protocol has to be applied some simple polynomial that would involve multiplying or summing multilinear extensions.

So for the triangle counting problem the polynomial $g$ would take the form

$$ g(X, Y, Z) = \widetilde{f_A}(X, Y) \cdot \widetilde{f_A}(Y, Z) \cdot \widetilde{f_A}(X, Z) $$

or in the case of GKR protocol $g$ would be

$$ f_{r_i}^{(i)}(b, c) = \widetilde{add}_i(r_i, b, c) (\widetilde{W}_{i+1}(b) + \widetilde{W}_{i+1}(c)) + \widetilde{mult}_i(r_i, b, c) (\widetilde{W}_{i+1}(b) \widetilde{W}_{i+1}(c)) $$

As you may see these polynomials are multiplications and additions of multilinear extensions of quite low degrees.

It would be nice to somehow generalize the sum-check protocol to be able to deal with any polynomial type and the above polynomials in particular. Thankfully Rust has all the tools needed to perform this kind of abstraction.

Abstracting Sum Check protocol over the type of polynomial

What actions does the Sum Check protocol perform with the given polynomial anyway?

  1. First the protocol needs to know the number of variables in the polynomial.

  2. Then obviously the Verifier of the protocol has to be able to evaluate the polynomial at a random point at the end of the last round.

  3. Recall, that on the $j$-th round Prover fixes all but one variables in $g$ turning it into a univariate polynomial $g_j(X_j)$ being $$ \sum_{(x_{j+1},...,x_{\nu}) \in \lbrace 0,1 \rbrace^{\nu - j}} g(r_i,...,r_{j-1},X_j,x_{j+1},...,x_{\nu}) $$ This suggests that the polynomial has to have a way to fix all but one variables.

  4. And finally there has to be a way to get the evaluations of this polynomial over the boolean hypercube. Summing these evaluations gives the value $c_1$ that Prover claims to be the true answer at the beginning of the first round.

These can be described as a trait in Rust like this:

pub trait SumCheckPolynomial<F: Field> {
    fn evaluate(&self, point: &[F]) -> Option<F>;

    fn to_univariate_at_point(
        &self,
        at: usize,
        point: &[F],
    ) -> Option<univariate::SparsePolynomial<F>>;

    fn num_vars(&self) -> usize;

    fn to_evaluations(&self) -> Vec<F>;
}

For now some methods in this trait return Option-al values in case some errors happen for instance, if the dimension of the point does not match the dimension of the polynomial.

Now, to reflect this change in the existing code of the Sum Check protocol that was implemented in the previous post all we have to do is implement the trait above for multivariate::SparsePolynomial:

impl<F: Field> SumCheckPolynomial<F> for multivariate::SparsePolynomial<F, SparseTerm> {
    ...
}

and change Prover and Verifier to be generic over this new trait:

pub struct Verifier<F: Field, P: SumCheckPolynomial<F>> {
    ...
}


pub struct Prover<F: Field, P: SumCheckPolynomial<F>> {
    ...
}

The rest of the changes are purely mechanical and can be found in this commit.

Sum Check Protocol to Count Triangles

For the triangle counting problem we are going to be dealing with the polynomial already mentioned above:

$$ g(X, Y, Z) = \widetilde{f_A}(X, Y) \cdot \widetilde{f_A}(Y, Z) \cdot \widetilde{f_A}(X, Z) $$

What is this polynomial? Well, suppose we are given a graph G, E - a set of edges in G and $A$ is the adjacency matrix of the graph $G$, i.e. $A_{i,j} = 1 \Leftrightarrow (i, j) \in F$. The matrix $A$ is used to count the number of triangles in this graph.

To do that matrix $A$ is viewed not as a matrix but rather as a function $f_A$ mapping $\lbrace 0, 1 \rbrace^{\log n} \times \lbrace 0, 1 \rbrace^{\log n} \rightarrow \lbrace 0, 1 \rbrace$ where indices $i,j$ are viewed as arrays of bits of $\log n$. Then the formula for counting the number of triangles would look like this:

$$ \Delta = \frac{1}{6} \sum_{ x,y,x \in \lbrace 0, 1 \rbrace ^{\log n}} f_A(x,y) \cdot f_A(y,z) \cdot f_A(x,z) $$

To define the polynomial that can be used in the Sum Check protocol the multilinear extension $\widetilde{f_A}$ of $f_A$ is used and so we arrive at the above polynomial $g(X, Y, Z)$.

For more detailed description of going from the $A$ to the polynomial
you should check out The Book, we are now going to get to implementing
this in Rust.

Since we are going to deal with the polynomial that is the multiplication of three evaluations of the same multilinear extension function over a different set of variables we can define our polynomial as follows:

pub struct G<F: Field> {
    f_a: DenseMultilinearExtension<F>,
}

Here we use a DenseMultilinearExtension from the ark_poly crate to implement the $\widetilde{f_A}$.

First lets implement the constructor that would allow to go from the adjacency matrix $A$ to our multilinear extension.

Being generic over any iterable structure of bool values:

impl<F: Field> G<F> {
    pub fn new_adj_matrix<M>(num_vars: usize, matrix: M) -> Self
    where
        M: IntoIterator<Item = bool>,
    {
        let g = DenseMultilinearExtension::from_evaluations_vec(
            num_vars,
            matrix
                .into_iter()
                .map(|b| if b { F::one() } else { F::zero() })
                .collect(),
        );

        Self { f_a: g }
    }
}

NOTE arkworks uses assertions within implementations of methods, the from_evaluations_vec constructor will panic if the dimensions of input parameters do not match.

Next, to use this polynomial with the new generic version of the Sum Check protocol defined above the SumCheckPolynomial trait for it. As I am going to implement it as close to The Book as possible it is going to get quite ugly.

First, the evaluate method. Recall, that the polynomial has the input of the form $(x_1,\dots,x_n,y_1,\dots,y_n,z_1,\dots,z_n)$ and our multilinear extension f_a has to be evaluated three times over respective parts of this input:

impl<F: FftField> SumCheckPolynomial<F> for G<F> {
    fn evaluate(&self, point: &[F]) -> Option<F> {
        assert!(point.len() == (self.f_a.num_vars() / 2) * 3);

        let mut x_z = point[..self.f_a.num_vars() / 2].to_owned();
        x_z.extend_from_slice(&point[self.f_a.num_vars()..]);
        // X, Y
        Some(
            self.f_a.evaluate(&point[..self.f_a.num_vars()])? *
            // Y, Z
            self.f_a.evaluate(&point[self.f_a.num_vars() / 2 .. ])?
            * self.f_a.evaluate(&x_z)?,
        )
    }

Here the cases of $(x, y)$ and $(y,z)$ are trivial since they can be simply taken as subslices of the input, however the $(x, z)$ case involves constructing the new vector since we are limited by how the contiguous memory works and I haven't found a better in-place solution yet. If DenseMultilinearExtension::evaluate API would be able to take something iterable instead of a slice, one could just chain to iterators of subslices.

Next, one trivial method to compute the number of variables based on the number of variables of the underlying DenseMultilinearExtension:

    fn num_vars(&self) -> usize {
        (self.f_a.num_vars() / 2) * 3
    }

The to_evaluations method is going to be a bit more tricky. What we have at hand is the f_a multilinear extension. The trait MultilinearExtension provides us with a to_evaluations method that "returns a list of evaluations over the domain, which is the boolean hypercube". But the domain of $f_A$ is $(X, Y)$ and in our case the domain of $g$ is $(X, Y, Z)$. It is clear that at each point in its domain $g$ is a multiplication of exactly three values from the vector of evaluations of $f_A$ over $f_A$'s domain, in other words these said three values are just values from f_a.to_evaluations() vector. But how do we compute the needed values to correctly index into this vector? Some good old bitmasking and bit shifting of course:

    fn to_evaluations(&self) -> Vec<F> {
        let x_size = self.f_a.num_vars() / 2;
        let bit_mask_most_significant = ((usize::MAX << x_size) ^ usize::MAX) << x_size;

        let mut res = vec![F::zero(); 2usize.pow(self.num_vars() as u32)];
        let evaluations = self.f_a.to_evaluations();

        for (x_y, evaluation) in evaluations.iter().enumerate() {
            for z in 0..2usize.pow(x_size as u32) {
                let f_x_y = evaluation;
                let f_y_z_idx = ((x_y << x_size) & bit_mask_most_significant) | z;
                let f_y_z = evaluations[f_y_z_idx];

                let f_x_z_idx = (x_y & bit_mask_most_significant) | z;
                let f_x_z = evaluations[f_x_z_idx];

                let f_x_y_z = (*f_x_y * f_y_z) * f_x_z;

                res[(x_y << x_size) | z] = f_x_y_z;
            }
        }

        res
    }

And now the last one: going to a univariate polynomial by fixing all variables but one.

Let's think about it:

Suppose we are given a point $(x_1,\dots,x_n,y_1,\dots,y_n,z_1,\dots,z_n)$ and the index $i$ of the variable that has to remain un-fixed so to speak.

Index $i$ lands into either $X$, $Y$ or $Z$:

And as such it would affect two out of three factors $\widetilde{f_A}$ in our polynomial. The third factor can be just evaluated at a constant point since it is not affected by this change.

Going to a univariate polynomial

The rust code that handles the case when $i$ lands in $X$:

    fn to_univariate_at_point(&self, at: usize, point: &[F]) -> Option<SparsePolynomial<F>> {
        let x_y = &point[..self.f_a.num_vars()];
        let y_z = &point[self.f_a.num_vars() / 2..];
        let mut x_z = point[..self.f_a.num_vars() / 2].to_owned();
        x_z.extend_from_slice(&point[self.f_a.num_vars()..]);

        match at / (self.f_a.num_vars() / 2) {
            0 => {
                let f_y_z = self.f_a.evaluate(y_z).unwrap();

                let a = self.g_to_univariate_at(at, x_y);
                let b = self.g_to_univariate_at(at, &x_z);
                Some((&(&a * &b) * f_y_z).into())
            }
        /* Cases for 1 and 2 are similar */

Again, we have to go use the same trick we used above to use a new vector to construct a value for $(X, Z)$. Another thing is that the helper g_to_univariate_at helper is used here that allows to fix all but one variables in f_a and turn it into a univariate polynomial.

How could this helper be implemented? Let's look at all the building blocks that can go into this implementation:

DenseMultilinearExtension implements the MultilinearExtension trait that among other things allows to fix a number of variables:

Reduce the number of variables of self by fixing the partial_point.len() variables at partial_point.

So we can go to the fewer variable polynomial by fixing the first $i$ variables.

It also implements a method that allows us to swapping variables (see the replace method).

How can we use these two? In a very unelegant way:

Suppose we have a point where we want to go to univariate poly at $i$:

$$ (x_1,\dots,x_{i-1},x_i,x_{i+1},\dots,x_n,y_1,\dots,y_n) $$

  1. Use fix_variables at point $(x_1,\dots,x_{i-1})$, now we have: $$ (x_i,x_{i+1},\dots,x_n,y_1,\dots,y_n) $$

  2. Use replace method to swap $x_i$ and $y_n$, to arrive at: $$ (y_n,x_{i+1},\dots,x_n,y_1,\dots,y_{n-1},x_i) $$

  3. Again, use fix_variables at point $(y_n)$: $$ (x_{i+1},\dots,x_n,y_1,\dots,y_{n-1},x_i) $$

  4. And finally use fix_variables again, this time at point $(x_{i+1},\dots,y_{n-1})$: $$ (x_i) $$

The last two steps could be squashed into a single one if we could swap to points in a vector, but we are using a slice instead of a vector.

At the end of these series of steps we will end up with a univariate DenseMultilinearExtension over variable $i$.

Now the only thing separating us from the finalizing implementation of the needed trait is turning this univariate extension into a univariate univariate::SparsePolynomial, since, you know, we are using Rust to typesafely guarantee that the polynomial is actually univariate.

What we know is that $g$ is a polynomial with degree at most $2$ in each variable so this univariate extension is an extension polynomial of degree of at most $2$. Quadratic polynomials can be described with at most $3$ coefficients. So we actually can interpolate our quadratic SparsePolynomial from evaluations of the extension. Luckily, ark-poly gives us exactly the machinery to do this with the types GeneralEvaluationDomain and Evaluations that stores a univariate polynomial in the evaluations form. Then we just call the interpolate method on the latter to get to the needed polynomial.

Here is the code for all the steps above that implements the said helper:

impl<F: FftField> G<F> {
    fn g_to_univariate_at(&self, at: usize, point: &[F]) -> DensePolynomial<F> {
        let mut fixed_1 = self.f_a.fix_variables(&point[..at]);

        if at != self.f_a.num_vars() - 1 {
            fixed_1.relabel_inplace(0, fixed_1.num_vars() - 1, 1);
            fixed_1 = fixed_1.fix_variables(&[point[point.len() - 1]]);
            let fixed_2 = fixed_1.fix_variables(&point[at + 1..point.len() - 1]);
            fixed_1 = fixed_2;
        }

        let domain = GeneralEvaluationDomain::new(3).unwrap();

        let evaluations = domain
            .elements()
            .map(|e| fixed_1.evaluate(&[e]).unwrap())
            .collect();

        let evaluations = Evaluations::from_vec_and_domain(evaluations, domain);

        evaluations.interpolate()
    }
}

Testing the code.

To test the code we need to use the $\mathbb{F}_p$ where $p$ is a prime such that $p \geq 6n^3$. We are going to write a simple test for the $4 \times 4$ matrix and the smallest $p$ satisfying this bound is $389$.

So lets construct the field, the adjacency matrix and kick off the Sum Check protocol the same way we did in testing in the previous post:

#[test]
fn test_simple_matrix() {
    #[derive(MontConfig)]
    #[modulus = "389"]
    #[generator = "2"]
    struct FrConfig;

    type Fp389 = Fp64<MontBackend<FrConfig, 1>>;

    let rng = &mut test_rng();

    let adj_matrix = vec![
        vec![false, true, true, false],
        vec![true, false, true, false],
        vec![true, true, false, false],
        vec![false, false, false, false],
    ];

    let g: G<Fp389> =
        G::new_adj_matrix(adj_matrix.len(), adj_matrix.iter().flatten().map(|b| *b));

    let num_vars = g.num_vars();
    let mut prover = Prover::new(g.clone());
    let c_1 = prover.c_1();
    let mut r_j = Fp389::one();
    let mut verifier = Verifier::new(num_vars, c_1, g);

    for j in 0..num_vars {
        let g_j = prover.round(r_j, j).unwrap();
        let verifier_res = verifier.round(g_j, rng).unwrap();
        match verifier_res {
            VerifierRoundResult::JthRound(r) => {
                r_j = r;
            }
            VerifierRoundResult::FinalRound(res) => {
                assert!(res);
                return;
            }
        }
    }

    panic!("should have returned on FinalRound from verifier");
}

and when we run it hopefully we will see that it succeeds:

running 1 test
test tests::test_simple_matrix ... ok

test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.02s

Final Notes

The code for the implementation as always may be found in the repo. The existing approach to Sum Check protocol is not quite efficient yet it has some obvious places to optimize as well as implement the optimizations The Book talks about. However recall that Premature optimization is the root of all evil and we were more focusing here on generalizing the protocol to lay the foundation for the upcoming GKR implementation. For any questions, suggestions or fixes just create an issue in the repo. Thanks and see you in the next post.

Changelog

26-08-2022