Towards Secure and Private AI: A Framework for Decentralized Inference

Despite the significant strides made in AI security in recent years, the potency of attacks has surged. Central to AI security is the pivotal task of delineating the threat model and understanding how adversaries target the inference process. Adversaries can exploit various vulnerabilities within the inference systems of foundational models, employing tactics tailored to different scenarios. In decentralized AI inference environments, one threat model emerges, where computing nodes may behave deceitfully, compromising the integrity of aggregated results. It becomes imperative to establish mechanisms wherein each node can verify its adherence to agreed-upon protocols without compromising the confidentiality of its model. This necessitates enabling nodes to provide proofs of honest execution to the central server or the public while safeguarding the confidentiality of their respective models. Thus, ensuring both the integrity of inference processes and the privacy of model architectures becomes paramount in the realm of AI security.

In the realm of defending against adversarial attacks, a plethora of meticulously crafted countermeasures exist to safeguard systems. One such strategy, particularly pertinent in decentralized inference settings, involves the implementation of mechanisms where central servers solicit proof of computing execution from each node. This process is orchestrated with precision to ensure that while the proof is furnished, the node’s confidential model parameters remain undisclosed. Subsequently, the central servers meticulously scrutinize the proofs submitted, thereby enabling them to discern the reliability of each node. The efficacy of this approach hinges upon the successful verification of the proof, serving as a litmus test for the trustworthiness of the node in question.

The challenge intensifies when fast inference of secure foundation models is required. Given the scale of big data and the models involved, foundation models inherently exhibit slow inference speeds. It is widely acknowledged that incorporating security measures further exacerbates this slowdown in AI models. For instance, the fastest-known zero-knowledge proof algorithm currently takes as long as 15 minutes to generate proof for a single token in the output [12]. Despite substantial efforts to expedite the inference process, it is imperative to ensure the security of foundation model inference without compromising efficiency.

Decentralized Inference. We explore the challenge of decentralized inference for foundation models, which offers numerous benefits. With the advent of 5G technology and improved internet latency, personalized devices such as mobile phones can now participate in crowdsourcing machine learning models. This approach enhances device utilization and eliminates the need for data communication between nodes and a central server, thereby safeguarding data privacy. In the context of Machine Learning as a Service (MLaaS), models remain the private property of individual nodes, allowing them to offer model inference services via APIs without sharing model weights and checkpoints. As foundation models grow in size—such as Meta’s LLaMA-3.1 with 405 billion parameters [4], and future models expected to reach trillion-level sizes—it becomes impractical to load entire models on a single node or server. Additionally, technologies like blockchain are now equipped to support the decentralized inference of foundation models on-chain, further facilitating this approach.

Model Decomposition. A core assumption in decentralized AI inference is that the model used in the inference system can be divided into several parts, each managed by a separate party, or computing node. Each computing node performs its assigned computations and sends the results to a central server. For instance, with a foundation model like LLaMA-3, the model can be decomposed layer-wise, enabling sequential inference by different nodes. Alternatively, decomposing the model width-wise allows for parallel inference across multiple nodes. The method of decomposition depends entirely on the application scenarios, and we consider both approaches in our products.

In this paper, we focus on decentralized inference of foundation models as described above. In this framework, each computing node (i.e. prover) owns a foundation model or a part of the foundation model with a publicly known architecture, while the model weights are proprietary. We make a semi-dishonest assumption on the central server (i.e. verifier): the central server is honest in aggregating results from each computing node and accurately reports the verification result, but the central server tries to glean additional information about parts of the foundation model at computing nodes. However, some computing nodes might be dishonest, potentially deviating from the agreed protocol by substituting their model with an alternative or outputting random data. We assume that the majority of nodes are honest, but acknowledge that dishonest nodes can collude. These adversarial nodes can arbitrarily alter their results, provided their behavior remains undetected by the central server.

We use the technique of zero-knowledge proofs to guarantee that each party is honest about his or her execution of inference of foundation models. Zero-knowledge proof serves as a fundamental technique and underpins the architecture of blockchain. In this cryptographic concept, two entities are involved: the prover and the verifier. The prover’s objective is to demonstrate the successful execution of a protocol without disclosing confidential information, termed as the ’witness’. This witness encompasses sensitive data like model weights or private information that the prover wishes to keep undisclosed to the verifier. Often, the protocol is depicted as a circuit, where certain components remain hidden within the witness.

Commitment, Proof, and Verification. The process of zero-knowledge proof involves three essential steps. Firstly, the prover commits to the witness data, such as model parameters, ensuring its integrity by encrypting it before transmitting it to the verifier. Once sent, the content remains unchanged, with the verifier gaining access only to the encrypted version. In practice, the (generalized Pedersen) commitment of any $d$-dimensional tensor $\textbf{S}=(S_{1},S_{1},...,S_{d})\in\{0,1,...,|\mathbb{G}|-1\}^{d}$ (e.g., the model weights) is implemented as [16]:

where $\mathbb{G}$ is an elliptic curve group w.r.t. the finite field $\mathbb{F}$ consisting of points $(x,y)\in\mathbb{F}\times\mathbb{F}$ such that $y^{2}=x^{3}+ax+b$ for designated field elements $a$ and $b$, $r\sim\{0,1,...,|\mathbb{G}|-1\}$ is an integer (i.e., the element of the scalar field $\mathbb{F}$ of $\mathbb{G}$ with $|\mathbb{F}|=|\mathbb{G}|$) uniformly sampled from $\{0,1,...,|\mathbb{G}|-1\}$, $\textbf{g}=(g_{1},g_{2},...,g_{d})\in\mathbb{G}^{d}$ and $h\sim\mathbb{G}$ are uniformly and independently sampled from $\mathbb{G}$. In the second step, the prover executes the protocol and simultaneously generates a proof of execution in a finite field $\mathbb{F}$, which is then transmitted to the verifier, either interactively or non-interactively. Depending on the protocol, operations can involve arithmetic or non-arithmetic processes. Lastly, the verifier meticulously examines the proof to ensure the honest execution of the protocol by the prover, thereby validating the transaction or information under scrutiny.

Zero-knowledge proofs have many advantageous properties that form the foundation of blockchain and decentralized machine learning. These include:

Completeness: If the prover accurately executes the circuit, the proof will be validated (with probability 1).

Special Soundness: If the prover is dishonest in executing the circuit, the proof will fail (with high probability). A weaker property, special soundness, requires executing the protocol twice and being able to identify the witness.

Zero-Knowledge: The verifier gains no knowledge about the prover’s witness.

The above properties ensure that the proof will pass verification only if the prover is honest, while also allowing the prover to keep its secret or witness hidden from the verifier.

Interactive vs. Non-Interactive Proof Systems. Depending on whether the proof-verification process involves a single round or multiple rounds, zero-knowledge proofs can be classified as either interactive or non-interactive, respectively. Zero-knowledge proofs are naturally described as an interactive process, where the verifier sends a challenge (typically a random variable) to the prover, who then responds to the verifier. If the proof is valid, the verifier sends a new challenge in the next round, and the process repeats (see Figure 1). Given an input $x$, an interactive proof procedure works as follows:

1. 1.

   $P$ sends the first message $\alpha\leftarrow P(x)$.

2. 2.

   $V$ sends a challenge $\beta$.

3. 3.

   $P$ sends the second message $\gamma\leftarrow P(x,\alpha,\beta)$.

4. 4.

   $V$ decides to accept or reject according to an algorithm $V(x,\alpha,\beta,\gamma)$.

In the non-interactive case, the process can be simulated by having the prover generate his/her own challenge $\beta$. This is achieved by replacing the random variable $\beta$ in the challenge with a random oracle model or a hash function $H$ (such as SHA-256) that operates on all the messages that the prover has sent so far. This is also known as the Fiat-Shamir heuristic [17]. In particular, given an input $x$, a one-shot, non-interactive proof procedure works as follows:

1. 1.

   $P$ computes the first message $\alpha\leftarrow P(x)$.

2. 2.

   $P$ computes a challenge $\beta\leftarrow H(x,\alpha)$.

3. 3.

   $P$ computes the second message $\gamma\leftarrow P(x,\alpha,\beta)$.

4. 4.

   $P$ sends $\alpha$ and $\gamma$ to $V$.

5. 5.

   $V$ computes $\bar{\beta}\leftarrow H(x,\alpha)$ and decides to accept or reject according to an algorithm $V(x,\alpha,\bar{\beta},\gamma)$.

The Pedersen commitment (1) satisfies the binding property: once sent to the verifier, the opening information $(r,\mathbf{S})$ cannot be changed anymore by the prover. Upon initial inspection, the prover needs to send the witness $\mathbf{S}$ to the verifier to prove that he/she can “open” the commitment. Thus the zero-knowledge property of the commitment (1) may appear impossible. However, this skepticism is unfounded, largely owing to the homomorphic property of the Pedersen commitment (1): for two commitments $\textsf{Commit}(\textbf{S}_{1},r_{1})$ and $\textsf{Commit}(\textbf{S}_{2},r_{2})$ corresponding to tensors $\textbf{S}_{1}$ and $\textbf{S}_{2}$, respectively, we have $\textsf{Commit}(\textbf{S}_{1},r_{1})\cdot\textsf{Commit}(\textbf{S}_{2},r_{2})=\textsf{Commit}(\textbf{S}_{1}+\textbf{S}_{2},r_{1}+r_{2})$, responding to the commitment of tensor $\textbf{S}_{1}+\textbf{S}_{2}$. This property enables the prover to prove to the verifier that he/she can “open” the commitment without revealing the witness. This is achieved by letting the prover instead open the commitment of a linear transformation of the witness: $\mathbf{S}^{\prime}\leftarrow\mathbf{S}\cdot e+\mathbf{D}$, where $\mathbf{D}$ is a $d$-dimensional hiding vector picked by the prover and $e$ is a challenge (scalar) randomly sampled by the verifier. Hereby, the vector $\mathbf{D}$ is to hide the witness $\mathbf{S}$ as $\mathbf{S}\cdot e+\mathbf{D}$ looks random to the verifier. The existence of challenge $e$ guarantees special soundness as two runs of the procedure with challenges $e_{1}$ and $e_{2}$ with $e_{1}\not=e_{2}$ satisfy:

In Equation (2), we have $2d$ unknowns $\mathbf{S}$ and $\mathbf{D}$ and $2d$ equations. Thus, two accepting transcripts $(\mathbf{S},e_{1},\mathbf{D})$ and $(\mathbf{S},e_{2},\mathbf{D})$ will identify $\mathbf{S}$ and $\mathbf{D}$ by Gaussian elimination, thus achieving special soundness.

Algorithm 1 provides a complete procedure for the opening of the commitment. Let $\mathbf{y}$ be a publicly known vector to both parties. The algorithm enables the prover to prove to the verifier that he/she knows a witness $\mathbf{S}\in\mathbb{F}^{d}$ and a witness $t\in\mathbb{F}$ which are openings of the commitment (1) and the commitment of $t$:

such that $\langle\mathbf{S},\mathbf{y}\rangle=t$, where $h$ and $g$ are generators sampled randomly from an elliptic curve group. The algorithm is provably complete, has special soundness, and exhibits zero-knowledge properties [18]. To see this, the completeness and zero-knowledge are obvious by the inspection of the algorithm. To see the special soundness, for two runs of the algorithm with challenges $e_{1}$ and $e_{2}$ ($e_{1}\not=e_{2}$) from the verifier, by Lines 8-9 of Algorithm 1, we have

where $S_{1i}^{\prime}$ and $S_{2i}^{\prime}$ represent the $i$-th element of vector $\mathbf{S}_{1}^{\prime}$ and $\mathbf{S}_{2}^{\prime}$, respectively. Denote by

Dividing Equation (4a) by (4b), we have

Similarly, dividing Equation (4c) by (4d), we have

Thus, we have shown that $(\bar{\mathbf{S}},r_{\bar{S}})$ is an opening of the commitment $c_{S}$ and $(\langle\bar{\mathbf{S}},\mathbf{y}\rangle,r_{\bar{t}})$ is an opening of the commitment $c_{t}$, as desired.

Multilinear Extension. Zero-knowledge proof focuses on the finite field $\mathbb{F}$, rather than the real field $\mathbb{R}$ as in the floating-point calculations. Arithmetic operations consist of addition and multiplication. For these arithmetic operations, it is easy to use the Sum-Check protocol, in particular, the GKR protocol [19, 20], to implement zero-knowledge proofs. The basic idea in the Sum-Check protocol is to express a $d$-dimensional tensor $\mathbf{S}\in\mathbb{F}^{d}$ involved in the calculation as a multi-variable polynomial $\widetilde{\mathbf{S}}:\mathbb{F}^{\log_{2}d}\rightarrow\mathbb{F}$ via a transformation called multilinear extension [18]:

where $\mathbf{b}\in\{0,1\}^{\log_{2}d}$ refers to the binary index of tensor $\mathbf{S}$, and $\widetilde{\beta}(\cdot,\cdot):\mathbb{F}^{\log_{2}d}\times\mathbb{F}^{\log_{2}d}\rightarrow\mathbb{F}$ is the (unique) Lagrangian interpolation polynomial:

such that for any $\mathbf{b}_{1},\mathbf{b}_{2}\in\{0,1\}^{\log_{2}d}$, the interpolation property holds true:

That is, the polynomial $\widetilde{\textbf{S}}$ is the Lagrange interpolation of the tensor S on the binary indices: $\widetilde{\textbf{S}}$ is the unique multilinear polynomial over $\mathbb{F}$ such that $\widetilde{\textbf{S}}(\mathbf{u})=\textbf{S}(\mathbf{u})$ for all $\mathbf{u}\in\{0,1\}^{d}$. With the multilinear extension, we can write the verification of arithmetic operations between vectors equivalently as the verification of sum of multi-variable, low-degree polynomial $g$:

by considering the multi-linear extension of tensors.

Sum-Check/GKR Protocol. Algorithm 2 describes the Sum-Check protocol, a.k.a. the GKR protocol [19], for Equation (11). The protocol proceeds in $v$ rounds [18]. In the first round, the prover sends a polynomial $g_{1}(X_{1})$ and claims it to be

where we use the capital letter (e.g., $X_{1}$) to represent the argument of the polynomial. A key observation is that if the polynomial $g_{1}(X_{1})$ is as claimed in (12) and $H$ is as claimed in (11), then $H=g_{1}(0)+g_{1}(1)$. If so, the remaining proofs then proceed by proving Equation (12). By the Schwartz-Zippel Lemma, it suffices for the prover to prove that Equation (12) holds at a random point $r_{1}\sim\mathbb{F}$:

In the second round, the prover sends a polynomial $g_{2}(X_{2})$ and claims it to be

Observe that if the polynomial $g_{2}(X_{2})$ is as claimed in (14) and $g_{1}(r_{1})$ is as claimed in (13), then $g_{1}(r_{1})=g_{2}(0)+g_{2}(1)$. By the Schwartz-Zippel Lemma, it suffices for the prover to prove that Equation (14) holds at a random point $r_{2}\sim\mathbb{F}$:

In the third round, the prover sends a polynomial $g_{3}(X_{3})$ and claims it to be

The recursive argument then continues with the proof of

In the final round, all $\overset{\text{?}}{=}$’s can be proved w.h.p. (i.e. with failure probability at most $\frac{d}{|\mathbb{F}|}$, where $d$ is the degree of $g$) by the recursive argument if and only if $g_{v}(r_{v})=g(r_{1},r_{2},...,r_{v})$. The latter can be easily verified by the verifier via the opening of the commitment to $g$.

The procedure of the Sum-Check protocol is shown in Figure 2. The completeness is straightforward by the construction. The soundness follows from the recursive application of the Schwartz-Zippel Lemma for each round. The zero-knowledge property follows from the privacy guarantee of the Pedersen commitment.

Reducing Linear Layers to Sum-Check Protocol. Linear layers frequently appear in foundation models. Mathematically, linear layers can be written as matrix multiplication. Given $n\times n$ matrices $\mathbf{A}$ and $\mathbf{B}$ and we denote $\mathbf{A}\mathbf{B}$ as $\mathbf{C}$. One can interpret $\mathbf{A},\mathbf{B},\mathbf{C}$ as function $f_{A},f_{B},f_{C}:\{0,1\}^{\log_{2}n}\times\{0,1\}^{\log_{2}n}\rightarrow\mathbb{F}$:

where sequence $(i_{1},...,i_{\log_{2}n})$ and $(j_{1},...,j_{\log_{2}n})$ are the binary representations of $i$ and $j$, respectively. Let $\widetilde{f_{A}},\widetilde{f_{B}},\widetilde{f_{C}}:\mathbb{F}^{\log_{2}n}\times\mathbb{F}^{\log_{2}n}\rightarrow\mathbb{F}$ denote the multilinear extension of $f_{A},f_{B},f_{C}$. It is easy to check that

Equation (19) is in the form of (11). Thus, we can apply the Sum-Check protocol for verifying the execution of linear layers with $\log_{2}n$-variant, degree-2 polynomial $g(\mathbf{b})=\widetilde{f_{A}}(i_{1},...,i_{\log_{2}n},\mathbf{b})\cdot\widetilde{f_{B}}(\mathbf{b},j_{1},...,j_{\log_{2}n})$ and $H=\widetilde{f_{C}}(i_{1},...,i_{\log_{2}n},j_{1},...,j_{\log_{2}n})$ at a random point $(i_{1},...,i_{\log_{2}n},j_{1},...,j_{\log_{2}n})\in\mathbb{F}^{\log_{2}n\times\log_{2}n}$.

There are two major techniques for zero-knowledge proof of non-arithmetic operations: 1) bit decomposition and 2) lookup table.

Reducing ReLU Activation to Bit Decomposition. Bit decomposition is one of the most frequently used techniques for zero-knowledge proofs of non-arithmetic operations. Let $\mathbf{Z}$ denote the pre-activated tensor of ReLU, and let $\mathbf{A}$ denote the after-activated tensor. We now introduce the techniques appearing in [10]. In the ReLU activation, we have

where $\odot$ is the Hadamard product, $\text{abs}(z)$ is the element-wise absolute value of $\mathbf{Z}$, and

is applied element-wise to the tensor $\mathbf{Z}$. Denote by $z=(z_{0},z_{1},...,z_{Q})$ the Q-bit binary representation of $z$, where $Q=\log_{2}p$ (e.g., 128 or 256), $p$ is the (prime) order of the finite field $\mathbb{F}$, $z_{0}=\text{sign}(z)\in\{0,1\}$ represents the sign, and $(z_{1},...,z_{Q})=\text{abs}(z)$ refers to the Q-bit unsigned binary representation of $z$. To prove Equation (20) holds for all elements $z$, we only need to prove: 1) $(z_{0},z_{1},...,z_{Q})$ is binary representation such that all elements are in $\{0,1\}$: $0\overset{\text{?}}{=}z_{i}(z_{i}-1)$ for $i\in\{0,1,...,Q\}$; 2) $(z_{0},z_{1},...,z_{Q})$ can recover $z$ in the sense that

and 3) Equation (20) is correctly executed for element $a\in\mathbf{A}$ in the sense that

Therefore, with bit decomposition, we can reduce the non-arithmetic ReLU operation to the arithmetic operations in the proofs of 1) and 2) and call the Sum-Check protocol. We can also use the Schwartz-Zippel Lemma to merge all the proofs in 1) and 2) into a single one:

where $r\sim\mathbb{F}$ is a random sample from the finite field $\mathbb{F}$.

Reducing Non-Arithmetic Operations to Lookup Tables. A lookup table is another technique for zero-knowledge proofs of non-arithmetic operations. It resolves the issue in the bit decomposition approach that the binary representation is lengthy when $p$ is large, so one non-arithmetic operation is represented by as many as $\log_{2}p$ bit-operations [18]. Hereby, the length of $\log_{2}p$ comes from the use of base-2 in the bit decomposition. In order to resolve the issue, one could consider a larger base $k$ and check the correction of base-$k$ decomposition by a table lookup. We now introduce the techniques appearing in [12]. In particular, both the prover and the verifier parties maintain a shared table: $\mathcal{T}:=\{(\mathbf{T}_{\mathcal{X}},f(\mathbf{T}_{\mathcal{X}}))\}\in\mathbb{F}^{n_{2}\times(d_{1}+d_{2})}$ for a non-arithmetic operation $f$ (e.g., the Softmax operation). To reduce the non-arithmetic operation $f$ to a lookup table argument, the prover needs to convince the verifier that the $d_{1}$-dim input vector $\mathbf{X}$ and the $d_{2}$-dim output vector $\mathbf{Y}$ satisfy $\sum_{i=0}^{d_{1}-1}r^{i}\mathbf{X}_{i}+\sum_{i=0}^{d_{2}-1}r^{i+d_{1}}\mathbf{Y}_{i}\subset\sum_{i}\textsf{Col}_{i}(\mathcal{T})r^{i}$ for a random $r\sim\mathbb{F}$, where $\textsf{Col}_{i}(\mathcal{T})$ refers to the $i$-th column of table $\mathcal{T}$. One can further reduce the proof of this “subset” relations to a Sum-Check proof, by observing that: $\mathbf{S}\subset\mathbf{T}$ for $\mathbf{S}\in\mathbb{F}^{n_{1}}$ and $\mathbf{T}\in\mathbb{F}^{n_{2}}$, if and only if there exists $\mathbf{e}=(e_{0},e_{1},...,e_{n_{2}-1})\in\mathbb{F}^{n_{2}}$ such that the following two polynomials over $X$ is identical:

By taking the logarithmic derivative at both sides, the following two rational functions are identical:

The prover sets and commits to $e_{i}=|\{j:\mathbf{S}_{j}=\mathbf{T}_{i}\}|$. Equation (26) can be proved by the Sum-Check protocol with a random variable $X\sim\mathbb{F}$, as all the operations here are arithmetic [12].

By integrating all components, one can construct a zero-knowledge proof system for foundation models like Meta’s LLaMA series. Foundation models comprise transformer layers followed by MLP layers. For arithmetic operations, such as the linear layers in transformers and MLPs, the Sum-Check protocol can be effectively employed to generate proofs. For non-arithmetic operations, like Softmax and SwiGLU activation, a lookup table can be utilized. For piece-wise linear activation, such as ReLU, bit decomposition can be used. Table 1 lists the zero-knowledge proof techniques for different operations in foundation models. Note that each operation requires carefully designed techniques for efficiency. To the best of our knowledge, zkLLM [12] is the first work that introduces zero-knowledge proofs to foundation models with tens of billions of parameters such as OPT [21] and LLaMA-2 [22].

We build a zero-knowledge system, zkDPS, for decentralized large language models. zkDPS is built upon zkLLM [12] and zkDL [10]. zkLLM [12] and zkDL [10] are state-of-the-art CUDA implementation of zero-knowledge proofs of machine learning. However, speed remains the most significant pain point for zero-knowledge proofs of foundation models. For example, it takes zkLLM roughly 15 minutes to generate a proof for the inference procedure of LLaMA-2 13B with an input token length of 2,048 [12]. In comparison, vanilla inference of zkLLM takes less than 0.1 seconds per token. Despite significant efforts have been made to accelerate generic zero-knowledge proof frameworks such as zk-STARK and zk-SNARK [23], they are still not fast enough, even much slower than zkLLM for billion-scale models. Therefore, the acceleration of zero-knowledge proofs is urgent for real-world applications.

zkDPS proposes the following ideas to speed up proof generation beyond zkLLM [12] and zkDL [10]. One idea could be to leverage the “compressible property” of LLMs by algorithmic innovation. It is well-known that LLMs are in essence a compression of information and data from the real world. Thus one could consider using a small draft model to “guess” the output of the original LLM, similar to speculative sampling. As the draft model is small (typically 20x-50x smaller than the original LLM [24]), one could expect to experience a shorter time for proof generation of the execution of the draft model. Experimentally, the top-1 guessing accuracy of the draft model is as high as 82% in EAGLE [24, 25], a state-of-the-art speculative sampling method. Therefore, this acceleration technique will not degrade the inference performance of LLMs too much.

Another idea is to leverage the hardware acceleration of modern GPUs. zkLLMs has already used CUDA implementation to parallelize the computations in the proof generation. For example, the CUDA kernels such as Pedersen commitment, elliptic curve, and lookup tables have been re-implemented or adapted in zkLLMs [12, 10], making it the fastest implementation of zero-knowledge large language models at the submission of this paper. We are optimizing those kernels by engineering methods such as Tensor Parallelism and KV cache.