facebook / winterfell

A STARK prover and verifier for arbitrary computations

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

GKR-LogUp: additional required API changes

plafer opened this issue · comments

The API changes in #271 to support GKR-LogUp were insufficient. As described in #257, GKR-LogUp relies on 2 auxiliary columns l (the Lagrange kernel) and s. As of #271, the prover currently doesn't have enough information to build the s column, and the Air doesn't have enough information to evaluate the s constraints.

GKR-LogUp auxiliary columns recap

The GKR circuit proof leaves the following claims to be proven by the Air: $\tilde{f_0}(r), ..., \widetilde{f_{m-1}}(r)$, where $\tilde{f_i}$ is the MLE of the i'th main trace column $f_i$, and $r$ is a random point (referred to as "Lagrange kernel random elements" in the code). At the time of writing, in the Miden VM's LogUp GKR implementation, these claims correspond to the FinalOpeningClaim.openings field.

To prove these $m$ claims, we first combine them into a single claim by taking a random linear combination. Hence, the claim to prove becomes:

$$ \lambda_0 \tilde{f_0}(r) + \dots + \lambda_1 \widetilde{f_{m-1}}(r) = \sum_{i=0}^{m-1} \lambda_i \tilde{f_i}(r) $$

This expression can be rearranged:

$$ \begin{align} \sum_{i=0}^{m-1} \lambda_i \tilde{f_i}(r) &= \sum_{i=0}^{m-1} \lambda_i \sum_{x=0}^{n-1} \mathrm{eq}(x, r) f_i(x) \\ &= \sum_{x=0}^{n-1} \mathrm{eq}(x, r) \sum_{i=0}^{m-1} \lambda_i f_i(x) \end{align} $$

Notice that this is an inner product between $$\mathrm{eq}(x, r)$$ and $$\sum_{i=0}^{m-1} \lambda_i f_i(x).$$ The columns $l$ and $s$ will implement this inner product. Specifically,

  • The $l$ column stores the Lagrange kernel. That is, the value at row $x$ is $\mathrm{eq}(x, r)$ (where $x$ is the bit decomposition of the row index), and $r$ is the above mentioned Lagrange kernel randomness.
  • The $s$ column is a running sum where row $x$ stores the inner product from row $0$ to row $x$. Specifically, $$s(x) = s(x-1) + l(x) \sum_{i=0}^{m-1} \lambda_i f_i(x)$$

The constraints for $l$ were discussed and implemented in #271. The constraints for $s$ will tie everything together, ultimately ensuring that the prover claims for $\tilde{f_i}(r)$ indeed correspond to the MLE of $f_i$ evaluated at $r$. Specifically, the boundary constraints are

$$ \begin{align} s[0] &= l[0] \sum_{i=0}^{m-1} \lambda_i f_i[0], \\ s[last] &= \sum_{i=0}^{m-1} \lambda_i \tilde{f_i}(r), \end{align} $$

where $\tilde{f_i}(r)$ are the above-mentioned prover claims left from the GKR proof. The transition constraint is:

$$ s[x+1] = s[x] + l[x+1] \sum_{i=0}^{m-1} \lambda_i f_i[x+1], \quad \forall x \in {0, \dots, n-2}, $$

where $n$ is the trace length.

Required API changes

There are currently 2 issues with the current API.

  1. the prover is currently missing access to the $\lambda_i$ values to build the aux trace, while
  2. the Air is currently missing access to both the $\lambda_i$ values and the $\tilde{f_i}(r)$ openings to evaluate the constraints.

Below are my suggestions on how to concretely achieve this (specific naming choices subject to change).

Define a new type GkrRandElements as

struct GkrRandElements<E> {
  lagrange: LagrangeRandElements,
  lambdas: Vec<E>
}

Consequently, AuxRandElements is modified to include this new type:

struct AuxRandElements<E> {
    rand_elements: Vec<E>,
    gkr: Option<GkrRandElements<E>>,
}

Then, Prover::build_aux_trace(main_trace, aux_rand_elements) now has all the information it needs to build the auxiliary trace, solving issue 1.

Then, Prover::generate_gkr_proof() now returns GkrRandElements instead of only LagrangeKernelRandElements. Consequently, Prover::generate_gkr_proof() the prover will also need to commit to commit to openings $\tilde{f_i}(r)$, and sample the $\lambda_i$ values. Analogously, GkrVerifier::verify() also needs to return GkrRandElements.

As for the changes in Air, I suggest we add 2 methods analogous to Air::evaluate_aux_transition() and Air::get_aux_assertions(), that are only called when a GkrProof is present, and pass in the $\lambda_i$ values, as well as the $\tilde{f_i}(r)$ openings (boundary constraints only). This solves issue 2.

Note that unlike the Lagrange kernel constraints, the constraints for the s column can be defined outside of winterfell, as they only require the typical EvaluationFrame.

Thank you! The issue and the PR make sense and as I mentioned in the PR review, we can merge it after a couple of minor comments are addressed. I am wondering, however, if we should try to do a bigger refactoring around this as I think some of the assumptions we made initially, may not make as much sense now.

First, we assumed that the Lagrange Kernel (lk) column and the s column would be built by the user (i.e., outside of Winterfell). It now seems like lk column and maybe s column as well can be handled internally by Winterfell.

Second, we currently treat GKR and LK as two separate things. For example, we can set lagrange_kernel_aux_column_idx in the AirContext without specifying GkrProof and GkrVerifier associated types.

Addressing the first one should be relatively simple:

  1. Instead AirContext::lagrange_kernel_aux_column_idx we could have something like AirContext::enable_gkr which would be a boolean flag.
  2. When enable_gkr is set to true, the prover would first make a call to self.build_aux_trace (as it does now), and then would build lk and s columns and append them to the aux_trace. In a way, Prover::build_aux_trace() would go back to being responsible for building aux traces only for user-defined aux columns.
  3. One the constraint side, enable_gkr flag would also automatically enable lk and s column constraint evaluation. For lk column this is what we do already, and one open question for me is whether we can do the same for the s column (seems like so - but want to confirm).

For the second point, my thinking is not as clear yet. But a few preliminary thoughts:

  1. It may be good to get rid of Air::GkrVerifier associated type in favor of adding Air::verify_gkr_proof() function. This function would be called by the verifier whenever AirContext::enable_gkr is set to true.
  2. I also wonder if it may be possible to replace Air::verify_gkr_proof() with something more basic that is useful for both the prover and the verifier. Ideally, the user would specify something like GkrCircuit and then the prover would be able to generate the GkrProof for it and the verifier would know how to verify this proof.

I agree. It no longer makes sense to have a "Lagrange kernel"-centric API.

First, we assumed that the Lagrange Kernel (lk) column and the s column would be built by the user (i.e., outside of Winterfell). It now seems like lk column and maybe s column as well can be handled internally by Winterfell.

+++. The algorithm to generate these columns is always the same, so best handled internally rather than copy/pasted by users.

One the constraint side, enable_gkr flag would also automatically enable lk and s column constraint evaluation. For lk column this is what we do already, and one open question for me is whether we can do the same for the s column (seems like so - but want to confirm).

Certainly.

For the second point, my thinking is not as clear yet. But a few preliminary thoughts:

I agree with the general idea in this subsection. I believe this is what you're describing already, but I would see winterfell handling as much of LogUp-GKR as possible (i.e. define GkrProof struct, as much of the circuit as possible, etc). The user only needs to define the input layer of the LogUp-GKR circuit (best way to do that is TBD); the rest can be handled internally. We already provide some machinery for the Lagrange kernel and "s" columns, but these have little use without the entire LogUp-GKR circuit proving & verifying (so we might as well provide that too).