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 sumcheck 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?

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

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.

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_{j1},X_j,x_{j+1},...,x_{\nu}) $$ This suggests that the polynomial has to have a way to fix all but one variables.

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 inplace 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 unfixed so to speak.
Index $i$ lands into either $X$, $Y$ or $Z$:

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

$(x_1,\dots,x_n,y_1,\dots,y_{i1n},y_{in},y_{i+1n},\dots,y_n,z_1,\dots,z_n)$

$(x_1,\dots,x_n,y_1,\dots,y_n,z_1,\dots,z_{i12n},z_{i2n},z_{i+12n},\dots, z_n)$
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 atpartial_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_{i1},x_i,x_{i+1},\dots,x_n,y_1,\dots,y_n) $$

Use
fix_variables
at point $(x_1,\dots,x_{i1})$, now we have: $$ (x_i,x_{i+1},\dots,x_n,y_1,\dots,y_n) $$ 
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_{n1},x_i) $$ 
Again, use
fix_variables
at point $(y_n)$: $$ (x_{i+1},\dots,x_n,y_1,\dots,y_{n1},x_i) $$ 
And finally use
fix_variables
again, this time at point $(x_{i+1},\dots,y_{n1})$: $$ (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, arkpoly
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
26082022