Scaling up Trustless Neural Network Inference with Zero Knowledge Proofs

foreign [Music] so this is a bit of a book report about what happened when I tried to convince the machine learning people that zero knowledge proofs were very relevant to them so in that line it's this is Joint work with a bunch of machine learning professors okay so the motivation we came to this problem with is that right now today a lot of the least trusted parties in society are running machine learning models that have broad societal impacts now you may think that blockchain is a very low trust space but I bet that people trust Twitter like even less and I don't think people trust credit scores very much either nonetheless these things are changing a lot of the way that we interact in society a second thread is that machine learning as a service has been growing in popularity especially recently with the rise of large Foundation models so in this Paradigm you know you're maybe making a slide and you want a cool image so you'll send hey openai I want an image of teddy bears working on AI research underwater an opening I will send you back this beautiful image to put in this slide um now the core problem in both of these use cases is that well as the user you have no insight into what the model Runner is actually doing you have no idea if an open AI actually ran you know Dolly 2 or they just have a farm of digital artists that are really exploited okay so we want some notion of trust and verification in this type of model inference okay so if you ask a computer scientist they'll tell you that there are actually many approaches to this Beyond ZK okay so we had to learn about them and the first that they'll bring up is multi-party computation so in this case a number of separate parties will perform the machine learning operation together and of course because they're collaborating this requires that the parties are online simultaneously and there's a one of n trust model for privacy and correct execution that means that you need to trust that only one of the end parties is actually behaving honestly so it's possible for NPC to provide both privacy and validity but unfortunately it requires pretty high interaction and high bandwidth and at least at present there's a really high compute overhead if you want to ensure validity so as a result we looked into the literature and it looks like MPC can only handle relatively small models today another technique that's pretty popular is homomorphic encryption the idea here is you take your input data encrypt it and send it to some server that runs a machine learning model on the encrypted data without knowing about it so in this case homomorphic encryption is only targeted at privacy and not at guaranteeing validity the benefits are that you don't need to interact with the server Beyond sending your encrypted input but the downside is that there's a really massive compute overhead for NPC I said there was a big compute overhead but trust me this is really massive and so as a consequence MPC today as far as I know can only handle quite tiny models I mean recognizing digits okay so now we come to ZK so what is ZK adding that MPC and homomorphic encryption don't bring to the table here so with ZK while the server is doing inference it can generate a proof of valid execution and with that proof anyone can verify that the output was correctly generated so in this case ZK can ensure validity that is that the model was correctly executed and it can ensure a weaker notion of privacy than MPC and homomorphic encryption that is for ZK the model runner needs to know everything about the input whereas that's not necessary in NPC and homomorphic encryption but once you have the proof the input can be private from the rest of the world another upside is that there's no interaction required and a downside is that there's still a pretty massive computation overhead over the cost of inference so in this talk I'm going to talk about a method we developed to scale this the application of ZK by reducing the compute overhead from pretty small models to relatively large models okay so let me first situate you in the task and what I mean by small and big so a lot of Prior work in ZK for neural network inference focused on these two Benchmark tasks mnist and c410 the task in mnist is you get a black and white digit and you have to recognize which of those digits it is the post office uses something like this in cfar there are 32 by 32 pixels of 10 classes of things like car automobile horse boat and you have to choose for each image which of those 10 glasses applies so in our work we scale things up to approach the much larger imagenet data set so in this case we have 1000 classes image that happens to contain many breeds of dog so actually there are 200 breeds of dog and at 1 000 classes the images are much higher resolution they're 224 by 224 and this is really the first standard large Benchmark data set for this image classification task okay so here's what we can do we chose a model called mobilenet V2 which was developed originally to be run in low resource environments such as mobile phones for a range of input resolutions we're able to completely snark the execution of the model um and as you can see the proving time for a single input is still relatively large okay so even for the smallest input resolution of 96 by 96 we take over two minutes but we are able to scale to the full image of that resolution in about 20 minutes so this is using the Halo 2 back end with the Original IPA proving system as you can see the proof size is relatively small and the compete verification time thankfully is not as large as the proving time although perhaps still quite large okay so let me in the rest of the talk let me talk a little bit about how we did this and what we can do with it so when you think about snarking execution of a neural network we came at it by dividing the problem into three pieces the first piece is it turns out that the notion of difficulty of execution between neural networks run on gpus and snarking things in a zero case circuit is quite different so we tried to select the best architecture for this task the second is we had to devise a pretty optimal way of arithmeticizing this neural network inference operation into a ZK circuit and last we selected a proving system largely because we were most familiar with Halo 2. so that's not a very principal choice but one thing that I wanted to note is we decided to go with a general purpose proving system instead of a proving system devised specifically for neural networks and the reason we wanted to do that was to plug into the existing tooling ecosystem around Halo 2 and other similar general purpose proving systems okay so to dive in a little bit more on what type of architectures are easy to do in ZK the first thing to note is that it's very hard to deal with floating Point numbers within a ZK circuit since every variable in a ZK circuit is under the hood a large Prime field element so that essentially forces us to work with quantized models that is models that were specifically designed to have all of their weights be 8-bit integers for now the level of quantization in machine learning generally gives you a trade-off between accuracy and how few bits you've quantized the model into and in ZK this translates to a trade-off between accuracy and proving cost so we did a first pass at this among just off-the-shelf models and it seems like models optimized for low resource environments particularly on the edge are quite good at this and so we chose to start moblenet V2 you can download this model from the tensorflow website that's the second thing we did to make our search over architectures a bit simpler is we wrote a transpiler from the tensorflow Lite model format to Halo 2 circuits and then we applied all of our optimizations to the individual building blocks of this transpiler so this allowed us to handle the many weird low resource neural networks without massive programming overhead finally in the arithmetization back end we use the planckish arithmetization for Halo 2 and so we actually only used two types of features of Halo 2 we handle all the linear layers via custom Gates and we use the lookup tables to essentially build a lookup table of all all non-linearities and finally there's a small subtlety that we have to readjust the fixed point in the quantization and we're able to do some optimization by shoving that into the lookup table okay let me now finish by just giving you a general sense of one application that we can do with this um so there's sort of four settings as Jason alluded to so where the you can make the model either public or private you can make the data either public or private and you can combinatorially combine these so we have a couple ideas in each quadrant so if your model is private and your data is public maybe you're trying to sell your model and you want to demonstrate to the buyer that your model is any good without just revealing the model altogether if your data and model are both public and we think that this could be useful for on-chain verification as Json alluded to and finally if your data is private then in both the private and public model format we think that this could be useful for conducting an audit for for example in legal Discovery without forcing revelation of the entire data set that you hold okay so just to give one sample back of the envelope calculation of how much it would cost to use our snark to do verified machine learning model accuracy we made a very basic protocol uh where a model perspective buyer asked a model seller to verify the accuracy of a model on a randomly chosen test set so in this case you want to randomly sample from the test set find the accuracy on the random sample and from that you're going to get some statistical estimates of your overall model accuracy um so we ran the numbers and it turns out that if you want to have um you know within you have want to know the accuracy within five percent at a 95 confidence interval you need to sample 600 times and that will cost about ninety dollars to verify if you want to be within one percent that's going to cost you a little over two thousand dollars uh with our current implementation and just to give you a sense of whether that could be reasonable uh people generally are willing to spend in the high five figures maybe low six figures just to acquire data to train their model so maybe it's worthwhile to pay a couple thousand dollars to verify that what you're buying is legitimate all right so just to summarize uh we constructed the first snark we think that can scale to imagenet level and for this we had to choose the correctly quantized mile model write a compiler from tensorflow Lite to a Halo 2 plonkish arithmetization and then we ran some benchmarks on whether this sort of technique is reasonable for some concrete applications of verified inference so going forward we're excited to explore some more applications of verified inference and also try to scale this to different types of models particularly Transformers thanks and I don't know if there's any time for questions thank you thank you for one it's gonna be oh and what do you know it's our next presenter that's so smooth I'll just hand you this mic and you can keep it wonderful um so one question you mentioned a couple of other organizations you did can you reveal some optimizations you think you could do in the future that's a very hard question since if there is an optimization we thought we could do we generally just did it uh one thing I didn't mention is we tried to reduce the number of lookup arguments by sharing them across layers um in gen this actually modifies the computation slightly so there's a little trade-off between the machine learning model accuracy and the proving speed