ZK12: Towards ZK-Friendly Machine Learning Models - Alan Li

Transcript

Hi, this is Alan from Brevity. Thank you very much for having me. I would like to share some of our recent work on efficient ZKML. What's ZK friendly ML models? Hope you enjoy. So there's no question that ML is being widely used. And there are many real-world applications now that heavily depend on dependent on ML models. Like autonomous driving, health care, co-pilot, and etc. Such exponential application comes with exploding size of the models.

So in the year of 2012 when AlexNet comes out, it only has 60 million parameters. And now 10 years later, when it comes to palm, this number becomes 540 billion. So basically the size has blown up by several orders of magnitudes. Another problem that has been brought by the wide application of ML is the problem of up showing model integrity, especially in sensitive applications in health care, finance, laws, and so forth.

We need to make sure that a specific output is obtained by executing a specific model over a specific So this is where the ZKML comes into play. There has been ZKML frameworks that could prove the model integrity of decision trees, CNNs, and LLMs. However, there is still a large gap of ZKML from being practical. So taking a medium-sized Llama 2 for example, if we are going to be approving the model integrity of a llama 2 with 3 billion parameters, it will take us about 15 minutes to

prove the integrity of the inference of a single token. This is already pretty efficient compared to what it was 2 to 3 years ago, but it is still not feasible. Most existing work focuses on ZK sites. So, basically, we develop new efficient ZKML proof protocols for existing ML models. And our observation reveals that this is not enough. Not only knowledge in the ZK, but also wisdom in ML are needed to make ZKML great again.

So, here what we mean by wisdom in ML. As we have mentioned, the model sizes have ever exploding, causing a lot of problems and issues in efficiency and pressures on computational resources. A series of work in ML called quantization and pruning resolves such issues. Machine learning models usually comes with the model parameters in the form of floating points of size 16 bits or even 32 bits. Quantization is the procedure of reducing the precision of model parameters from floating point to integer of 8 bits or even less.

While models could be very large, if we plot the number of parameters versus their values, we can see that most of the model parameters have absolute values that is very close to zero. Which means these parameters contribute little to the final outcome of the model. By using pruning, we are able to remove unnecessary or the low impact parameters by simply setting them to zero. And this also saves the computation. If we combine quantization and pruning, we are able to achieve more than 10 times of reduction in the number of sizes of parameters.

But the caveat is that heavy quantization and pruning will lead to severe loss in model performances. As we have mentioned, the quantization have different levels, and we can quantize the original model with with floating points parameters of say 16 bits or 32 bits to integers of eight bits, four bits, two bits, and finally ternary and binary networks.

The binary network have parameters quantized to plus or minus one. It has a long history that dates back to 1990. It achieves the extreme of quantization, so all the model parameters are only one bit. But it comes with a sacrificing model performances. On the contrary, ternary networks are proposed in in in the year of 2016, which is after the binary networks. Which weights are quantized to plus or minus one or zero, which is why it is called ternary.

It is also achieves the limit It It actually almost achieves the limit of quantization and it saves a huge amount of computation, but without loss of model performances. This might be counterintuitive, which is why I want to emphasize again. While ternary networks are extreme case of quantization, they are simply more efficient without sacrificing model performances. We take a middle-sized Lama for example.

The perplexity of language modeling of now ternary network is even better than its full precision counterpart with with significant reduction in latency and memory footprints. Such ternary network is not only useful for enhancing the model inference efficiency, but also very ZK-friendly in a sense that it saves a huge amount of ZK-related computations All right, we can take a linear layer as a showcase.

For a linear layer in the form of W * X = Y, where W is the full precision matrix, then the Y is going to be computed as a sum of field products. However, if we are going to be to be using a ternary linear layer, then each entry in W is quantized to either +1 or -1 or 0. Then the computation of Y is transformed into sparse additions and subtractions and excludes all the multiplications.

All right, now everything seems perfect. First thing is that operations amounts are reduced. Ternary weights naturally naturally brings a sparsity in model parameters, reducing the total number of required computations. Second, operations types are transformed. All multiplications operations are transformed into additions and subtractions, which is vastly reducing the computational burdens. Last but not least, it is a gain without pain.

Efficiency gain comes without loss of model capacity, i.e. model performances are maintained. So, the question becomes that is it mean the ternary network saves ZKML? The answer is no. Or to be more precise, it is not enough. As we have mentioned before, we now have wisdom from ML world. But we need knowledge from ZK as well. Now the models are sparse and with ternary entries. So, there's a call for sparsity and a ternary aware ZKML protocols.

To leverage the sparsity and the ternary structure of the linear layer, we propose SPASON, a variant of sumcheck. It is able to prove a linear layer in a time that is proportional to the number of non-zeros in ternary matrices. So, if the ternary matrix is sparse, which is expected, then using SPASON will save a lot of time in proving the model integrity. The intuition here is that we could save computation on those zero entries, as well as plus and minus one entries, when doing bookkeeping of the sumcheck.

There's a lot of technical details behind. If you are interested, you could refer to our paper online. Now we have a special series models, ternary networks, that are ZK friendly from wisdom of ML. We also designed special protocol from ZK perspective to better leverage the special proper properties of the ternary networks. So, here we show some very preliminary results by combining wisdom of ML and the knowledge of ZK.

We are able to boost the efficiency of proving a linear layer by three orders of magnitudes compared to a naive application of GKR on a linear layer of normal ML model. It is also worth mentioning that such technique can be generalized to handle not only linear layers, but also general layers. Basically, we generalize GQR to leverage the sparsity of ternary networks, and the resulting protocol called SpaGQR is able to prove one general layer in time that is proportional to the number of non-zeros in ternary matrices.

This work is still ongoing, and we expect to release more in-depth results in the future. All right, that's all for our presentation. We are Brevity. Thank you.