ZK14 - Jolt Atlas: Verifiable Inference via Lookup Arguments in Zero Knowledge - Antoine Douchet

Transcript

Today we'll introduce you to George Atlas, the framework that we built at I see me lab and are still building to provide verifiable agents. If my clicker works. I will use my keys. So, today we will first mention our motivations. Then I will give you some architectural ideas of how we design things.

Then I will provide you with preliminaries to give you a few examples of operators that we prove with George Atlas. And I will conclude with benchmarks and our future work. So, the motivation is of course led by the huge increase in usage of artificial intelligence in the current years in the past years. Uh, so it is now involved in high stakes usage as medical or even taking financial decisions. Uh, on the contrary, it's still treated as a black box where we give it a prompt and we just hope that the result is high quality.

So, we have a gap between the expectations that we have about the AI that we want it to take good decisions about our finances and all and the actual trust that we can allow it. So, this is where we introduce George Atlas to in fact give the artificial intelligence auditability. Uh, we had a first step as George, which is given by the name and we will see how George program doesn't really fit machine learning.

So, first of all, George is interesting because it allows us by being a virtual machine to abstract the complexity of writing circuits which can lead to a lot of errors. It also has under hood the sum check protocol, which proves everything from the instruction correct correctness until also the memory checking. So, this is a unique primitive and it makes it highly parallel parallelizable and all. Uh also, the philosophy of the Jolt virtual machine, which is look up singularity, is very interesting for us because it The idea is that we can prove any instruction given we have a good look up table for that.

And this is in fact interesting for us because machine learning is full of weird operators such as softmax, hyperbolic tangent, and such. However, Jolt is not really good for machine learning because you need to commit to the full trace, and when you have millions, if not billions, of parameters, this is highly prohibitive. Also, even simple instructions as tensor multiplication or tensor sum, when you compile that to Jolt, you have an array of individual element-wise sum over the elements, each element of the tensors.

This is not even mentioning complex operators such as matrix multiplications. So, we took a step back and looked at the GKR protocol, which has a few interesting points for us. First of all, you only send the output layer, so the output of your model to the verifier, and you give him a commitment to the inputs. The For the rest, the verifier just needs the structure of the intermediate layers, but doesn't need to know the exact values of that.

Also, it is great for tensor wide operations because you step from the output layer to the next layer by a unique iteration of the sum check protocol, and then you repeat that sequentially to eventually get to a claim about the input. And if you verify the claim about the input, it verifies the claim about the output of the model. However, as we can see, it has expensive sum checks which have are over a lot a lot of elements and with selector polynomials. So, we refined the model to get to layer uniform circuits where each layer has the same operation over the whole layer element twice.

This gives us in our view, uh layer tensor equivalence. So, now we will not really treat layer, but rather tensors as they are used in machine learning models. Uh also, we greatly simplified the sum check protocol and the issue that is remaining is that we still only can only handle sequential dependency. So, if we have one input tensor, we cannot fit it to different operators. It can only be used once. So, we we moved to a directed acyclic graph view of the things. And now, as we can see the two main consequences is that a sum check can now derive to two different claims.

And also, a tensor, so a layer, can have two claims coming from two different sum checks. Uh however, we know a protocol and we use a protocol to reduce those two claims into one. So, we can just keep in mind that we are uh iteratively going from the output layer to the input through sum checks. The only issue that remains now is that we are still only performing uh any model as long as it is only composed of addition and multiplications. Uh so, we just take for example an encoder block, which is really common in machine learning, and try to abstract away from the addition and multiplication and use that on our claim reduction flow.

So, all we see that it fits quite nicely. We only need to define exactly what are those sum checks that appear between each tensor. So, before stepping and seeing how actually we are proving any kind of machine learning operators, I will give you a few preliminaries. So, the sum check protocol allows to reduce a claim about the output of a tensor into a claim about just one evaluation of this polynomial. It is iterative and at each round the verifier samples a random element to give it to the prover.

So, for all notations, we will use for each operator A B C as the tensors. We will also use them to represent the multivariate polynomial that represents its values. A and B will be our inputs, C the output. As we can see on the multivariate polynomial, to take one index to index over the tensor, we will just consider the binary representation of the index that we want to represent. And we will use multilinear extensions for distance distance amplifications.

So, now I introduce you to sum check singularity. We are bouncing off of lookup singularity concept in in Jolt and consider that in fact we can prove any operator given we have the a good sum check. We will split into three main categories, element-wise operators we where the same operation is done on each element of the of the inputs. Reduction where we are computing each output given several inputs, several elements of the input also matrix multiplication softmax.

And indexing where we are merely reorganizing the input elements into another ordering. So, first I will give you a easy example, the multiplication to give you the grasp of what is interesting for us and what we need to prove, how do we prove it? What we want to prove is that for whole the whole hypercube, we have the above the equation that the element of the output equals the corresponding element of both inputs multiplied one to the other.

So, the first intuition that we could have is to use directly the Schwartz-Zippel lemma and just sample a random element evaluation point and verify whether when we evaluate the output to this element, it equals the evaluation of the inputs multiplied together. Uh however, this will not work because we cannot we do not have an equality of multilinear polynomials given the right hand side is not multilinear. So, the solution is to introduce P of X. Uh we can see it has the same notation as when we are trying to uh get the multilinear extension of any multivariate polynomial.

Uh given this construction, we know that P of X agrees with the product of A of X * B of X on all points of the hypercube and on top of that it is now multilinear. So, we just need now to prove that C of Z for a random uh Z given by the verifier equals the the polynomial evaluated at Z. And this is where we introduce the sum check protocol. We will prove this exact sum equals C of Z. So, here T equals C of Z uh by using the sum check. And now we are reducing a claim about the output C of Z into claims about the input. As we can see below, everything there can be can can be just computed by uh computing the claim of A

at R and be at all. So, now we are switching to another example. Here we are using Jolt's uh lookup table uh Jolt's way of proving lookup tables. So, we want to prove that every input of C, well, every element of C equals to the ReLU operator applied to its input. To do so, we will define a one-hot encoding of A such that for each element it will equal I only if uh this correspond the index of the vector, the one-hot vector, corresponds to the evaluation of A at this point. So, we

have an example here where if A of X equals 6, we will see that the only value that is set to one of these one-hot encoding is the value at index one. We will use that to create the polynomial P of X as we can see below. And again, under the example, evaluating this polynomial consists of just computing the dot product of our one-hot encoding and a table that contains all evaluations of the ReLU table at those indexes. And as we can see, it nicely just uh reduces to ReLU of 6.

Um since we are using Jolt's shoot algorithm to prove to prove the lookup table, we need to prove that it has prefix-suffix decomposition, which is the most important property for lookup tables in Sholt. And we are hence constructing the ReLU table as just a check on the sign bit, which is BS below. Uh and if the sign bit is uh set, so is equal to one, then it means the number is negative, so we are just remove slicing anything else of the number, and we just output zero as expected.

So, the protocol become Well, the sum check protocol applied to that is just verifying that T equals the below the above sum. And on top of that, we will need to verify the correct one-hot encoding of our input A. Which is done with other algorithm and virtual polynomials to follow also the notation given by George. Now, another example that we will give is for matrix multiplication. Here, we want to prove that for all elements of the output, it is the dot product of the A and B input just for the same dimension.

Here, we can notice that C is already multilinear in I and G, so we can directly feed that to the polynomial P. Uh so, that makes things easier for us. We can just directly apply the sum check protocol to this polynomial. Uh we can also notice something that is, in my opinion, interesting here is that with all work, it is as expensive merely to prove a matrix multiplication that it is to just prove a simple multiplication between two matrices.

Now, the last example will be for indexing indexing operators, which can be quite a challenge for us because usually we would think of it as um memory readdressing, and we would like to do some memory checking to assert that every element of the input exists in the output. However, we have nice tricks for each of those operators, which allow us to highly reduce the cost of proving them. For example, for move axis, where we are just switching the two dimensions of a polynomial or any dimensions, uh we can see that we want to prove that for all elements C of IG equals A of GI.

So, we are just indeed switching the two parts of the evaluation points. And again, C is already multilinear, so we can fit it directly to our polynomial P. And there, something even more interesting is that our polynomial P doesn't even involve any sum. So, we can abstract away from the sum check protocol, and we simply have to propagate the claim from the output to the input with basically no cost. So, all we do is we just divide the evaluation pointed to its two parts, invert them, and then propagate that for the input claim.

So, obviously, we still have to commit to the inputs of our model. And for that, Jolt uses Dory, and we choose to rather use Hyper KZG. So, in fact, the two the main differences between Dory and Hyper KZG is that Dory is sublinear. The generator set is subli- sublinear is rather than KZG, which is linear. Uh, but obviously, on the other hand, you have higher overheads. So, the estimations are that the advantage threshold for Dory it lies between 2 to the 30, so a number of elements to commit to that is 2 to the 30.

And in our setups, when we are proving GPT-2 or real world models real world models that we are able to prove now, we never reach this threshold, so this is why we decided to use Hyper KZG rather than Dory. So, now I will show you some evaluation and benchmarks of our model, of our product. So, we evaluated on GPT-2 with different sizes from the smaller one to the um the 125 million parameters one.

Uh so, we can see that for the 125 million parameters, we reach a 40-second proving time. We also used it to prove Qwen at 600 million parameters and had a 48-seconds proving time. When comparing to other works, we are the fastest and also these benchmarks should really be taken with a huge pinch of salt. Uh we are proving using a MacBook Pro M3 with 16 GB of RAM. And most of those benchmark comes come from highly powerful machines. So, we are the fastest and we are still not using the the most expensive machines.

So, future works that we are looking forward. On the first side, there is model hiding because currently the verifier, in order to verify all any model, needs to know the exact structure of the model. We can already hide the inputs, which is provided with Blindfold. But, we cannot we do not hide what operators are used and how they are linked together. So, model hiding would provide the ability to prove models for industry for which the exact nature of the model is a competitive advantage. So, they will be able to say, "I indeed used the model I told you."

And they will still be able to hide what's behind the model. On the other hand, there is lattice-based cryptography. This will allow us, obviously, to be post-quantum secure. But, also it will allow us some optimization and faster computation on the field elements because we will be able to use a 2 to the 64 uh sized field. Uh that would be all for today. Thank you for your attention. >> Thank you, Arthur.

Uh we do have time for questions. So, any questions? Yeah. >> Hey, uh thanks for the talk. Uh I'm not familiar with Hyper KZG, but uh does it also have a trusted setup like KZG? >> Uh I'm I'm sorry, I can't answer the the question, but I I believe we do not. >> I see. Okay. Uh thank you. What being based on? Yeah. >> How long does it take to run the inference without uh making the proof?

>> It would be like maybe a second or less. This is still a really high overhead. >> Yes. Okay, plenty of questions. That's great. We have the time. You'll be next. >> So, the arithmetic here is field arithmetic right? But the models they use floats, I guess. How do How does that work? So, like a quantized model? >> Yeah, I didn't talk about it in the talk, but indeed we are quantizing and we are not using float, but uh integers.

So, what we do, we use fixed point integers to represent the decimals. Which is easier for us to use look up tables on. >> So, if you do a multiplication, you need to reduce the bit size after How How does that work? >> Uh So, for multiplications, indeed we need to assess that the inputs are small enough, rather. And also, after the multiplication, we need to do rescaling because we would be adding the scale. So, this is uh ev- every multiplication is uh bending, like has a kind of division inside to also uh

do the rescaling. >> Thanks. >> Thank you. >> Um I can I can answer your question. Um so uh yeah, Hyper KZG still does a KZG commitment to the um the uh polynomial. It's equivalent to the MLE one. So, um yeah. The answer is that yes, it doesn't need to adjust the setup. >> Thank you. >> And we have one more up here. >> Um thanks for the talk. So, you showed a lot of operation like ReLU, matmul. Um how would it work for other types of operation like cos or uh tangent, stuff like that that are like not necessarily

like linear? >> So, for those operators, we since for example, I will just take the cosine operator. Um indeed this is not linear, so we cannot have a perfect suffix decomposition table. So, we will have a custom table, but we reduce its size by uh exploiting the fact that the cosine function, for example, is periodic. So, what we will do first is take the remainder of our input over 2 pi. I'm just simplifying, but we take the remainder over 2 pi, and then we use that uh output which is small enough into the cos lookup table.

>> Okay, and you're still able to run a sum check over it using the parity CD? >> Then we use the same technique uh uh yeah. So, here the perfect suffix decomposition allows us to have big sum check, so over 32 bits, uh which is prohibitly prohibitly expensive for custom the lookup tables. So, this is why we take the remainder to get a small number and then apply it to the cosine look-up table.

>> So, like you would have the sum cos and some like module operation or like remainder? >> So, in fact, we would first have here a division where we take the remainder. Then we take the one hot encoding of the remainder and then we use that into the look-up table. >> Awesome. Mercy. Mercy. >> Are there more questions? Yes. >> So, um obviously um using functions um that ideally are of a reals like um cos and so on. Um is the quantization that you use there for look-up tables uh is that similar to what you would use if you were um

uh running the model um outside the circuit or do you optimize that? >> Indeed, it's not the same quantization. So, we would have a different version of the model if we want to because if we run the model to get the output and then try to prove it with our model, we will have even if there are slides, we will have some differences which will prevent us to prove uh to have a valid proof. So, what we do is we have like this slightly different model that uses our quantization, but with results that are really identical.

>> Okay. >> We have time for one more. Yes, okay. >> I was also wondering how um do you then use like how do you then encode logs for like cuz you use those all the time in like attention and LLMs and stuff like that. >> Uh we in the models that we are proving, we do not require the log function. Uh otherwise I I believe I had worked on that and I'm pretty sure there is a look up uh a perfect suffix decomposition uh look up table for log that is possible to find. So, we can use the same technique as we do to prove log.

>> Okay. Thank you very much, everybody. Thank you Antoine for >> Thank you.