ZK13: ZKTorch: Efficiently Compiling ML Models to Zero-Knowledge Proof Protocols - Daniel Kang

Transcript

[Music] Uh, so I'm Daniel Kang. I'm a professor at UIU and today I'll be talking about ZK torch um and about how to compile machine learning inference to zero knowledge proofs via parallel accumulation. Uh but first unless you've been sleeping under a rock you will have known that AI is taking over the world. Uh there's new model releases every month. So just actually in the past uh six weeks or so 03 was announced along with the new version of Gemini 2.5 Pro which have shown amazing performance in different domains. Uh the issue with a lot of these models though is that they often change underneath us. So I'm going to give a concrete example. Um there have

been allegations that the 01 model release uh the system card accompanying the model release did not actually describe the release version of 01. And besides that, uh just two weeks ago, OpenAI released a new version of chat GPT uh which uh allegedly caused uh potentially negative uh uh consequences to some of their users who was who are using their platform. And so we have no idea what's really going on underneath the hood uh with these model providers. So uh this is one uh reason why we care about this. Uh I'm going to give you another example uh which is the case of medical AI. So let's say I'm uh u I I go to a doctor and I get an X-ray and I want to classify my uh X-ray using

a approved model. Let's say now you know here I'm showing the example of you might uh imagine that uh the correct model says one answer but then uh someone can hack the uh the system and actually just return a completely different answer. And this is obvious to see but there's actually lots of examples of uh instances where model deployers will deploy models as I said for example for chatgpt without actually understanding the consequences and so you might not actually know what model is being run on your own data and this is this can be benign it can also be malicious and as a model consumer I have no idea uh what uh um uh what model or or what are the guarantees of the model that was used to produce the result at

Now, of course, we're at uh ZK Summit, so this should be obvious to most of the people here that one of the ways that we can resolve some of these tensions is by using zero knowledge proofs of machine learning model inference. And so I won't go into the details of of this uh in too much uh too much depth, but basically uh every time you return a an inference from a model provider, you can attach a proof of computation saying that I ran say in this example an FDA approved model uh for this specific um on your specific data. Okay, so that's great. uh but there is this issue which is that unfortunately uh machine learning models are extremely expensive from a computational perspective and so now

we're stuck between this uh rock and a hard place. You can either have uh say generalpurpose technologies um but uh that can prove a wide v variety of models but they're unfortunately extremely slow uh or you can have these custom uh bespoke protocols which are unfortunately very specialized and don't actually uh uh prove a wide range of models. And in this talk I'll show how to uh start to bridge this gap. Okay. But first I'm going to talk about the two axes uh in the in the prior work uh before talking about ZK torch internals and then some benchmarks and applications if I have time left. Okay. So the first set of

techniques that uh people have used for ZKML broadly falls under compiling ML models to general purpose uh ZK uh circuits. So you take your favorite representation of a circuit or an air or something else. So you can imagine R1CS or um say Halo 2 style errors where you basically take the computation and machine learning model and then write a compiler from that to a general purpose circuit. Now this is obviously general because these circuits are universal. Uh but the issue is that you can have over a 100,000 times overheads with general purpose zero knowledge proofs. And so for small models like say linear regression uh or say small decision trees this can actually work just fine.

But the problem is as you get to larger and larger models which are becoming more prevalent in our lives, I'm sure most of you guys have used Chat GPT or some other form of LLM. Uh this 10,000 times overhead no longer becomes uh tenable. Okay. So then uh oh nice formatting issues. Uh well uh so the other uh approach that people have taken is you can develop custom protocols designed for specific models. And so for example uh here are two such examples in the literature zknn and zklm. Now one of the things I actually want to point out here is that this thing here GPTJ is actually an LLM but by default zklm does not support GBTJ and this is because of

the specific way the protocol is designed. Now this isn't an inherent limitation. you could extend it but the point I'm trying to make here is that if you want to prove a variety of different models uh the prior work is unfortunately inflexible and cannot support a wider range of models at least that's currently implemented. So key insight that we leverage to to making this more general and also higher performance is that ML models and this includes a wide range of uh different ML models that people actually use in practice uh consist of a small number of composable computations.

Uh I'm sure everyone here has heard of matrix multiplications. uh but there's also some other form other things like non different linearities that can then be uh that if we can prove each of these individually we can then put together to prove general purpose machine learning models and also a wide range of them. So to support this vision we built zk torch which is both extensible and high performance. The general way this works is that we have some specification of a uh of a model and currently we support onyx and we built a compiler which takes this specification and compiles it down to zk torch layers and then given these zk torch layers we have a transpiler to what we call our what we call basic blocks which I'll describe in detail but these are basically

cryptographic primitives and then we have the actual prover which then executes and creates the proof. uh uh as far as I'm aware uh ZK torch is the first uh proving system that supports all the MLF uh version 1.4.1 edge models um and it's about uh six times faster than general purpose uh zero knowledge proofs on LLMs and we also developed a new parallel accumulation scheme uh to be succinct and high performance. Okay, so now that I've talked about zero knowledge proofs and MLS in general and a little bit about the high level of ZK torch, I now want to talk about the internals of ZK torch and how we accomplish this. The first and most important uh

sorry the first thing that's really important is how we represent the uh intermediates in a machine learning model and these are generally called tensors. A tensor is simply a floatingoint uh n-dimensional array that is used uh that is used between uh layers of a model and also the inputs and the outputs. So roughly speaking let's say you uh let's say we use chat GPT as an example. Uh your text that you give to chat GBT is turned into actually a floatingoint uh tensor and then is passed through many layers which produces intermediate tensors before producing the output tensor which is then turned back into text. So these tensors are fundamental to to these machine learning models. So what we do in ZK torch is we

take the original tensor and then we quantize it and then we kg commit the tensor along the last axis. So here I'm showing a two dimensional uh array but you can also imagine n dimensional arrays where we just simply commit along the last axis. Uh one thing I want to point out actually is that you might want be worried about this quantization step. Well, it turns out that um a large fraction of machine model inference uh perhaps I don't it's hard to get the exact numbers but perhaps over 95% of it is actually done using quantized weights. So this is actually standard practice today outside of uh ZKML. Okay. So why do we use uh KG? Why do we use KG? Well, KG easil

easily supports lookups and arithmetic operations. uh homorphic addition is great if you want to add uh two cases you committed vectors that's very very cheap okay so given the tensors we want to now be able to perform operations over these tensors and so we selected uh uh 20 basic blocks which support the 61 layers that are necessary to prove the uh mlf uh v 4.1 models uh and the the we have a list here so they span from arithmetic and boolean operations So for example adding two two tensors as I mentioned uh table lookups for nonlinearities and scaling implemented using uh a modified version of CQ uh and this is for things like uh the value

nonlinearity and also for other operations like the softmax which are critical uh for modern machine learning and then we take uh some uh and then we uh also support matrix multiplication uh shape operations to change the shape of tensors and then also other things like sorting and max operations which are used for some of the more specialized models like retinaet. One of the key primitives that we use to implement these basic blocks is what we call a poly IOP. uh and the basic idea is that we know a set of witness polinomials and we have some uh know some relation f sorry relation r such that the r on these

polinomials is equal to uh zero. So for example uh this r here of this f for a boolean check basically to determine if all elements are zero or one can simply be f of one uh times 1 minus f of one and you should be able to convince yourself that if this relationship is satisfied then f the kg committed uh polinomial uh is all zero the coefficients are all zero or one. Okay, beyond the boolean check which was very simple as I mentioned uh our poly IOP enables many different operations. Uh so for example one of the uh uh uh checks that we use is we have an ordered check basically to check if

the if a vector is sorted in ascending or descending order doesn't doesn't really matter here. So let's say we have some some input. Uh what we can do is that we compute the difference between the uh uh the shifted difference uh using the polyop and then we simply uh check that all of the differences are positive and you should also be able to convince yourself that uh if the difference between consecutive elements in a vector are all positive then it is uh sorry I guess it's technically not negative uh then uh this should be an increasing vector.

Okay, so I talked a little bit about the basic blocks which is this uh component here. Now I want to talk about another component of our system which is the compiler. So uh operations that are efficient on the GPU are not necessarily efficient uh in uh uh ZK. So for example uh this is a nonlinearity that is widely used for LLMs. is called the gel nonlinearity and this is how it's implemented on the GPU. You have some input, you take it to some power, you multiply, add, multiply, 10h, add, multiply, multiply on the other side, and then get some output. So you can see here there's 1 2 3 4 5 6 7 8 different operations that are required

to implement this on the GPU. And this is because the GPU doesn't have this amazing thing which are these lookup tables. Uh and so instead of uh computing this, we actually take uh at the com at the onyx level uh replace using a um a pattern matching uh uh uh uh compiler uh these complex operations with operations that are simple uh for uh zk proving. So this is entirely at the onyx level. This is actually not this actually doesn't work at the prover level at this point. Okay. So I've talked about this part and this part. Uh now I want to talk about the transpiler. So one of the things about

uh uh ZK is that there are many ways to do the same computation. And here I'm just showing one simple thing, one simple uh way that you can do two uh uh two protocols for the same operation which is matrix multiplication. Now for matrix multiplication uh there's a protocol developed by um Ariel and uh Liam uh I believe uh called CQL and in CQIN you actually need to know one of the matrices but it's uh you can do a very efficient matrix multiplication but if you don't know the matrix multiplication we can use another protocol or variant of another protocol called Aurora to do the matrix multiplication and the basic idea is that you do a random linear a

randomized check uh and that's uh the rough intuition. Now deciding between these two, we need to know if one of these matrices is not known ahead of time and also depending on the size of the matrices and some other factors which one we want to use. And so here the transpiler decides which of the different uh um implementations of the same operation uh we select to to get high performance. Another uh thing another important class of operations that that the transpiler operates over are these shape operations. Uh here I'm just showing a a list of uh basically the decision tree for how we decide which shape operation to use. Uh the details

here aren't super important. Essentially there are many ways uh to do uh to to basically reshape or transpose or change the shape of a uh of a tensor. uh and which one we want to use depends uh quite a bit in terms of uh the generality versus the efficiency trade-off here. And so we want to use the most specialized version which is generally the highest performance version for shave operations. Okay, so I've talked about our compiler and our transpiler. Now I actually want to talk about the prover and specifically how we optimize certain uh parts of the proving stack to both uh obtain succinctness and also high performance.

Uh okay so um I want to talk about some prior work which is uh this idea of a proof accumulation. I think that we actually specifically use folding. Uh the general idea is that if you have uh two proofs uh generally speaking they're uh they let's just say they're the same operation for the time being. uh you can uh accumulate uh a1 into sorry p into a1 to get a proof uh a2 um that basically says that a1 and and p are correct and I have something that's a sync as uh after doing this this is a very standard technique in uh in zk why is this important well uh in zk torch we can produce thousands of of

proofs and of course if you produce thousands of proofs you're going to have a very very large proof size which is not desirable and in fact in some of our uh some large models uh I think we can actually produce tens of thousands of proofs so to resolve this we extend this prior work by mirror by uh Benfish and some others um I believe at Yale uh and by using accumulation we can reduce the proof size by over 2,000 times and this is again because uh ZK torch can produce thousands of small proofs depending on the uh specific operations and the emission learning All right. So, how do we extend mirror?

Well, okay. So, as far as I'm aware, there aren't really that many uh uh settings where you are producing literally tens of thousands of proofs. Uh and so, uh the way that this uh mirror does things in general is that they do this uh proof accumulation sequentially. And instead, we can basically do a Merkel accumulation, which is substantially more efficient if you're accumulating thousands and thousands of proofs. Okay, so so far I've talked about uh ZK torch internals. Now I want to move on to some benchmarks. So we evaluated all the models in the MLPF uh v1.4 edge version and the hardware that we used was this uh uh you know Intel processor with 64

threads and 4 terabytes of memory. The model spanned from LLMs to CNN's uh to assorted ones. And here we have RNNs and also this model called stable diffusion which is used for image generation. Here I'm showing the endto-end performance including the proving time, the verification time and the proof size. Uh as we can see uh the uh we can prove all these models. The proof times obviously vary. Stable diffusion for example is very very very large. Uh um and so it's obviously more more complex but as we can see here the verification time is generally relatively low. uh this can be reduced further using other tricks like uh recursive verifi uh verification and what I want to show

here also is that the proof sizes are also generally low for these for the size of these ML models in addition by using mirror we can actually produce smaller proofs so for example here I'm uh showing proof sizes for uh ResNet 50 for mystique and zk torch I believe actually a smaller these smaller models for zknn and pbcnnn and so zk torch has uh more than three times smaller proving sizes even using a larger CNN compared to prior work. Uh and then finally I want to point out that our parallel implementation of mirror improves proving time by six times uh reducing you know the prover for say GBTJ from about 8,000 seconds to about 1400

seconds and the verification is also substantially reduced as well. Um can I get a quick time check? 10 minutes. Great. Okay. So now I've talked about um uh zero groups in ML, some ZTO internals and benchmarks and I want to talk about some potential applications of uh ZKML. Okay. So the first thing I want to talk about is uh how generative AI works which is that you have some inputs and you have uh a prompt for these inputs.

For example, I'm showing here stable diffusion where we have the prompt of an astronaut riding a horse which then generates uh this uh this this image here. Uh, one of the fun facts or not a fun fact, one a fun anecdote is that I uh visited uh the folks who work on midjourney and they actually have uh um uh something in their office which actually shows the images that are generated live and you can actually see people playing around with different prompts especially for work when they want to get something that specifically uh targets what they have in mind. And so this process of generating going from a prompt to an image actually is a bit of an art. Now beyond uh image generation uh this is actually also true for LLMs. As it turns out you have to

often do this thing called prompt engineering to achieve high performance using LLMs in different circumstances. Uh so for example this image here was generated from this uh fairly long prompt. The prompts can actually even get long can get even longer than this. And let's say I wanted to generate this this image. Uh hopefully there's no one as offended by this. uh you know IO would have no idea how to go about uh doing this. And so there's there's a question of how do you actually uh uh uh generate these complex prompts for uh these images that you want to be able to uh create from these models. Okay, so this was as of you know a couple years ago, but nonetheless uh you know there are these prompt

engineering jobs that can pay up to $300,000 a year. So quite valuable using uh zk proof uh we can take a prompt and then actually uh uh have the output and then have a proof of the computation while hiding the input. And so this is one way that you can imagine actually uh proving to say an employer or someone else that you can actually uh generate high quality images or high quality prompts for LLMs uh while preserving your IP.

Okay. Uh right. So this is uh prompt engineering uh but doing it in a provable way using zero knowledge proofs. Okay. Beyond uh LLMs uh and image generation for high performance. I also want to talk about uh this idea of auditing ML systems. Well, okay. For first of all, why do we need why do we need audits? Well, there's many reasons. uh one is that there's this amazing line of work uh from Chaffy Goldwater at all uh that shows that you can actually plant undetectable backdoors in ML models uh which is quite scary which basically means that on specific inputs the ML model can uh you can make the ML model output something specific and it this is

undetectable from simply looking at the weights. Now beyond this, I showed earlier in the talk that there were allegations that the 01 model released by OpenAI, the model card accompanying it was not actually reflective of the actual model that was deployed in production. Okay, so uh one of the things that zero knowledge proofs enables is this idea of a trustless audit. And a trustless audit is exactly what it sounds like. I don't need to trust the trainer or the model provider. uh and I can actually be assured of some properties of this model. So one example is that I can prove that you ran an FDA regulated model. You can also do something like proving uh say my uh Twitter timeline is

not biased. Okay. So in order to do this uh we will going to do two steps. The first is that we're going to do first do a proof of training and in this proof of training we first commit to the training data in a way that um preserve preserves privacy and then we produce zk snarks of the training process itself. Now obviously today this is uh too expensive to do on say an LLM but you could imagine doing this in uh using fine-tuning in a costefficient manner uh or you can imagine also doing this on smaller models. Uh then uh if you want to do an audit, you can then write down another function which is a fun uh a function over the data and the model weights at the end of training

and you know that this model these model weights came from this data because you have this proof of training. Okay, so this audit function can be actually quite complex. uh you can even have an audit function where you run an LLM over the input data itself. So you can do quite complex things like saying uh you know does this input data refer to uh some sensitive property of a user that I don't want for legal reasons or something else of this form. So yeah, so this is you can uh actually prove that that the training data isn't biased uh that the model performs well and many basically any other computable property uh you can uh

uh prove using this procedure and I actually want to emphasize this one more time that this f can be arbitrary and in fact using a technique like uh a system like ZK torch you can in fact use LLMs as part of your audit which is uh um I think very powerful and allows you to start to do sort of these arbitrary computations over your input data. Okay, the final application I want to talk about uh is proof of personhood in a decentralized and trustless manner. So uh I'm sure everyone has seen deep fakes in some form or another. Uh this is a simple example from meta. Uh and you know here we have this don't go to this website by the way ftx compensation.com. Uh these deep defects

are used by state actors to spread misinformation tricking businesses into miswaring funds and also uh uh scamming uh people on the internet. So in order to fight deep fakes uh one thing you can do is use this uh uh uh this thing called a tested camera and a tested camera basically you put a trusted ex uh an enclave right on the on the device say the camera and you sign all the pixels that come off. So again um uh this is a well-known technique uh in the literature. What can you do with this?

Well, the first thing you can do is that you can edit images privately. Let's say we have this very sensitive credit card information here. Uh we can then uh just uh remove it and then prove that that uh edit was done honestly. Okay. What does this have to do with deep fakes? Well, one thing I can do uh all with this technology is I can then compose it and then do privacy preserving biometric authentication. So starting with an Nessa camera, I can prove that a photo again this is all private uh was uh originated from a real camera. I can then perform uh photo editing uh again using in a completely privacy preserving way using a zero knowledge proof and then actually run face ID again using a system like ZK

torch and in a completely privacy preserving way uh that I'm a real person while not actually revealing say who I actually am. And this can be combined with things like um uh zero knowledge proofs of uh let's say uh government uh sanctioned u identities and other things of this form to really uh increase trust while preserving anonymity. Uh beyond that I should gave some examples of audits. You can also for example uh audit say the Twitter algorithm or other other uh algorithms for for providing recommendations to your timeline. I'll just skip through that.

Okay. So I'm just going to conclude. Uh AI is becoming increasingly important as I believe everyone is aware. Uh and zero knowledge proofs provide a method of doing audits over both private models and private data which I think is very powerful. uh and we built a system called ZK torch which which is extensible and also high performance. Uh I was hoping to have this open sourced uh this week but unfortunately we're going to push it back till next week. So if you want to follow me on Twitter uh which will where we'll announce it. I also want to uh briefly announce that project veil is hiring. Uh if you want to chat with Jonah about that uh he's in the audience over there and I'm happy to take any questions at this point. Thanks everyone.

[Applause] for your transpiler. How many op codes are going from onyx? And when you do accumulation, are you paying for all those op codes per step like an IVC or are you doing something more like NIVC where you pay per opcode? Uh sorry, I think the the echo is unfortunately a bit difficult. So my my my understanding of the question is that are we doing accumulation at the onx layer level or at some other level? We're actually doing that at the basic block level and they're actually uh uh many uh onyx layers that actually get decomposed into many different basic blocks. Uh so it's a bit complex to determine what the relationship is. Um but uh we'll be open sourcing it. So if you want to do a deep dive uh when it comes out next week, happy to help with that. Cool. Thanks.

Hey, uh thank you for the talk. Um I noticed when you were talking about the benchmark stuff um you said them the specs you were running was 4 terabytes of RAM uh for generating the proofs. Obviously that's not feasible for a lot of individuals and even if like for example in your the example you gave like verify that this X-ray was was passed through this model even for like a hospital that that's just not feasible hardware. So I mean obviously we want to get to that point but in the meantime you would want to delegate that kind of computation to some kind of you know third party proer. Have you considered how like you would do privacy preservation in that context for obviously delegation? Uh yes the

question is that the hardware requirements are quite severe for for this work. Are there ways that we can do delegation or other techniques to reduce this? Uh there's a few things to say. The first thing is that purely in terms of memory you can simply cut this down by doing uh by proving say one layer at a time. Now obviously this would increase the proof size and you have to also increase the amount of computation because you need to prove that the uh the input the witness from the previous uh layer is actually the the correct one for the next layer. Uh but there are ways to reduce the uh total uh memory consumption. Now, in terms of delegation, um I haven't been keeping up to date with uh uh uh the uh the the mechanisms that allow you to do proving uh privacy

preserving delegation. But I'm pretty sure almost everything that we use is fairly standard and can be delegated uh if you want to do it in a privacy preserving uh way. Uh but I'll have to look into some of the specific details around that. Um we are working on next generation systems which uh already are giving us 5 10x speed ups over the results I'm showing here and we're hoping to uh reduce the the hardware requirements even more uh beyond that. Thanks. We got time for one more question before the break if anyone would like to ask one.

Oh yeah, go for it. Why do you choose mirror over the other? Okay, so the question is why do we choose mirror over sorry if I can just can you just repeat the question for the benefit of the audience? I've read mirror so I kind of understand but I don't think it's obvious like why did you choose mirror over the other folding scheme that are out there? Uh yeah so to to answer your question quite directly um we it it simply supported a bunch of things that we were doing already. Uh so it they fold uh CQ and CQL in uh and because that was already built into their protocol it was just a simple thing for us to extend. Um uh there are many different uh choices of uh protocol at the different layers and they actually result in very very different uh uh

performance characteristics. uh we didn't have a chance to look into others um and we're hoping to do this in the future but as I mentioned uh there are some techniques that we're already developing that can reduce the cost by five stacks or more and we're also looking to reduce the cost even more beyond that uh and this is going to require uh optimization across the whole stack. All right, let's all give Daniel a hand. Thank you. Thanks a lot.