Kernel-level proof mechanism of machine learning models

Original author: Zhiyong Fang

"How do you eat an elephant? One bite at a time."

In recent years, machine learning models have made tremendous advancements at an astonishing pace. With the enhancement of model capabilities, their complexity has also surged in tandem—today's advanced models often contain millions or even billions of parameters. To address these scale challenges, various zero-knowledge proof systems have emerged, which consistently strive to achieve a dynamic balance among proof time, verification time, and proof size.

Table 1: Exponential Growth of Model Parameter Scale

Although most of the current work in the field of zero-knowledge proofs focuses on optimizing the proof systems themselves, a key dimension is often overlooked—how to reasonably decompose large-scale models into smaller, more manageable submodules for proof. You might ask, why is this so important?

Now let's explain in detail:

The number of parameters in modern machine learning models often counts in the billions, and even without involving any cryptographic processing, they already occupy a very high amount of memory resources. In the context of Zero-Knowledge Proof (ZKP), this challenge is further amplified.

Each floating-point parameter must be converted to an element in the Arithmetic Field, and this conversion process itself will increase memory usage by approximately 5 to 10 times. Additionally, to accurately simulate floating-point operations in the Arithmetic Field, extra operational overhead must be introduced, usually around 5 times.

Overall, the memory requirements of the model may increase to 25 to 50 times the original scale. For example, a model with 1 billion 32-bit floating-point parameters may require 100 to 200 GB of memory just to store the converted parameters. Considering the overhead of intermediate computation values and the proof system itself, the overall memory usage easily exceeds the TB level.

Current mainstream proof systems, such as Groth 16 and Plonk, typically assume that all relevant data can be loaded into memory at the same time in their unoptimized implementations. While this assumption is technically feasible, it is extremely challenging under practical hardware conditions, significantly limiting the available proof computation resources.

Polyhedra's Solution: zkCuda

What is zkCuda?

As described in the zkCUDA Technical Documentation: Polyhedra's zkCUDA is a zero-knowledge computing environment aimed at high-performance circuit development, designed to enhance proof generation efficiency. Without sacrificing circuit expressiveness, zkCUDA can fully leverage the underlying provers and hardware parallel capabilities to achieve rapid ZK proof generation.

zkCUDA language is highly similar to CUDA in terms of syntax and semantics, making it very friendly to developers with existing CUDA experience, and it is implemented in Rust at the lower level, ensuring both safety and performance.

With zkCUDA, developers can:

• Quickly build high-performance ZK circuits;
• Efficiently schedule and utilize distributed hardware resources, such as GPU or cluster environments that support MPI, to achieve large-scale parallel computing.

Why choose zkCUDA?

zkCuda is a high-performance zero-knowledge computing framework inspired by GPU computing, capable of breaking down ultra-large-scale machine learning models into smaller, more manageable computing units (kernels), and achieving efficient control through a CUDA-like front-end language. This design brings the following key advantages:

1. Precise Matching Proof System Selection

zkCUDA supports fine-grained analysis of each computation kernel and matches it with the most suitable zero-knowledge proof system. For example:

• For highly parallel computing tasks, protocols such as GKR that are good at handling structured parallelism can be selected;
• For smaller or irregularly structured tasks, proof systems like Groth 16 that have low overhead in compact computation scenarios are more suitable.

By customizing the backend, zkCUDA can maximize the performance advantages of various ZK protocols.

2. More intelligent resource scheduling and parallel optimization

Different proof kernels have significantly different resource demands on CPU, memory, and I/O. zkCUDA can accurately assess the resource consumption of each task and intelligently schedule to maximize overall throughput.

More importantly, zkCUDA supports task distribution across heterogeneous computing platforms—including CPU, GPU, and FPGA—thus achieving optimal utilization of hardware resources and significantly enhancing system-level performance.

The Natural Fit of zkCuda and the GKR Protocol

Although zkCuda is designed as a general computing framework compatible with various zero-knowledge proof systems, it has a natural high compatibility with the GKR (Goldwasser-Kalai-Rothblum) protocol in terms of architecture.

In terms of architecture design, zkCUDA connects various sub-computation kernels by introducing a polynomial commitment mechanism, ensuring that all sub-computations operate based on consistent shared data. This mechanism is crucial for maintaining the system's completeness, but it also incurs significant computational costs.

In contrast, the GKR protocol offers a more efficient alternative pathway. Unlike traditional zero-knowledge systems that require each kernel to fully prove its internal constraints, GKR allows the verification of computational correctness to be recursively traced back from the kernel output to the input. This mechanism enables the propagation of correctness across kernels without needing to completely unfold the verification in each module. Its core idea is similar to gradient backpropagation in machine learning, where correctness claims are tracked and transmitted through a computational graph.

Although merging such "proof gradients" in multi-paths brings certain complexities, it is this mechanism that forms the deep cooperative foundation between zkCUDA and GKR. By aligning the structural characteristics in the machine learning training process, zkCUDA is expected to achieve tighter system integration and more efficient zero-knowledge proof generation in large model scenarios.

Preliminary Results and Future Directions

We have completed the initial development of the zkCuda framework and successfully tested it in multiple scenarios, including cryptographic hash functions such as Keccak and SHA-256, as well as small-scale machine learning models.

Looking ahead, we hope to further introduce a series of mature engineering techniques in modern machine learning training, such as memory-efficient scheduling and graph-level optimization. We believe that integrating these strategies into the zero-knowledge proof generation process will greatly enhance the performance boundaries and adaptability of the system.

This is just the starting point, zkCuda will continue to advance towards an efficient, highly scalable, and highly adaptable universal proof framework.

Original link