Private Neural Network using Zero-Knowledge Proofs during Inference.

Photo by Marija Zaric on Unsplash

Introduction

Privacy is critical in various applications, particularly in Neural Networks, where sensitive data is often involved. This project addresses a significant challenge in Neural Networks (NN) and Cryptography using modern Zero-Knowledge Proofs cryptographic protocols.

The primary objective is to preserve the Privacy of the NN scheme while still allowing the verification of computations inside the NN.

To achieve this, we leverage the recent advancements in this intersection of NN and Cryptography, particularly the optimized approach outlined in ZEN.

ZEN approach introduces a new kind of cryptography-friendly quantization that minimizes the accuracy loss typically associated with cryptographic operations on Integers over a finite field.

ZEN allows using privacy-preserving, secure, and efficient Neural Networks in various fields, including healthcare and face
recognition.

The ZEN model is a PyTorch implementation of this quantization scheme, which can be used to quantize the weights of a neural network. The ZEN model also provides a method to perform Inference on the quantized neural network in a privacy-preserving manner, using Zero-Knowledge Proofs.

What is a Zero-Knowledge Proof?

You can find a nice introduction to this concept in this video:

Let’s try to go into the details of how Zero-Knowledge Proofs (ZKP) work. In Cryptography, we traditionally use two characters, Alice and Bob, trying to share a secret or, in the ZKP case, trying to prove something.

Imagine two people getting divorced who are already separated and live in different cities, and of course, they don’t trust each other. They must decide what to do with some of the remaining items left in the house. So they agree on tossing a coin over the phone to decide who will keep the items.

This is an untrusted setting. How can Bob trust the outcome of the tossed coin by Alice? Or even better, how can we use Number theory to solve this trust problem?

The elegant solution comes from the hardness of problems based on Number Theory, like the Discrete Logarithm Problem.

The Discrete Logarithm problem states that if I have g^x mod p, where p is a prime number. finding x from here is a computationally hard problem when the prime p is big.

Using ZKP, we can prove, for example, that:

• A certain number is a quadratic residue,
• or that something is the encryption of zero
• or that two elements are encrypted using different bits.

all of which in a way that doesn’t reveal anything else.

The Quantization Problem

Quantization is the process of converting a constant value into a discrete value. The quantization process originates from the need to represent real numbers in a computer.

Since computers can only store a finite number of bits, it is impossible to represent all real numbers. Therefore, the quantization process converts a real number into a discrete value that can be represented in a computer.

The quantization process is also used in Neural Networks to convert the weights of a neural network into a discrete value that can be represented in a computer.

In general, most of the cryptographic operations occur on Integers over a finite field. For example, the Zksnark protocol uses elliptic curve cryptography (ECC) to perform cryptographic operations. The ECC is based on the discrete logarithm problem, a complex problem to solve. Therefore, the ECC performs the cryptographic operations on Integers over a finite field. However, the ECC could be more efficient when performing cryptographic operations on Integers over a finite field.

Convert a Neural Network into a Quantized Neural Network

We first implement a neural network using PyTorch. We use the CIFAR10 dataset, a widely used benchmark in computer vision research, consisting of 60,000 32x32 color images in 10 classes, with 6,000 images per class. The classes are airplane, automobile, bird, cat, deer, dog, frog, horse, ship, and truck, and the images are split into 50,000 training images and 10,000 test images.

This Neural Network we will call the Floating-point Neural Network. We use this Neural Network to train the CIFAR10 dataset.

class FloatingPointNet(torch.nn.Module):
    def __init__(self):
        super(FloatingPointNet, self).__init__()
        self.l1 = nn.Linear(784, 128, bias=False)
        self.act = nn.ReLU()
        self.l2 = nn.Linear(128, 10, bias=False)

    def forward(self, x):
        x = self.l1(x)
        x = self.act(x)
        x = self.l2(x)
        return x

Then we trained with nn.CrossEntropyLoss()

model = FloatingPointNet()
BS = 128
loss_function = nn.CrossEntropyLoss()
optim = torch.optim.Adam(model.parameters(), lr=0.001) # adam optimizer
losses, accuracies = [], []

for i in (t := trange(1000)):
    samp = np.random.randint(0, X_train.shape[0], size=(BS))
    X = torch.tensor(X_train[samp].reshape((-1, 28*28))).float()
    Y = torch.tensor(Y_train[samp]).long()
    optim.zero_grad()
    out = model(X)
    cat = torch.argmax(out, dim=1)
    accuracy = (cat == Y).float().mean()
    loss = loss_function(out, Y)
    loss.backward()
    optim.step()
    losses.append(loss.item())
    accuracies.append(accuracy.item())
    t.set_description("loss %.2f accuracy %.2f" % (loss.item(), accuracy.item()))

With this setting, we got an accuracy of ~94.9%.

And the most important for our project was saving the weights:

torch.save(model.state_dict(), "weights.pkl")

Quantized Neural Network with PyTorch

The next step is to convert the Floating-point Neural Network into a Quantized Neural Network. We use the ZEN model to quantize the weights of the Floating-point Neural Network. The ZEN model is a PyTorch implementation of the quantization scheme proposed in the ZEN paper.

ZK circuits for different layers

Zero Knowledge Proofs are very versatile and can be used to prove the correctness of some calculations and, when combined with SNARKs, to provide the Privacy layer.

In the case of proving correctness, Pedersen Commitments are very useful.

pub fn pedersen_commit_long_vector(
    x: &[u8],
    param: &PedersenParam,
    r: &PedersenRandomness,
) -> Vec<PedersenCommitment> {
    let len_per_commit = PERDERSON_WINDOW_NUM * PERDERSON_WINDOW_SIZE / 8; //for vec<u8> commitment
    let num_of_commit_needed = x.len() / len_per_commit + 1;
    let mut commit_res = Vec::new();
    for i in 0..num_of_commit_needed {
        let mut tmp = Vec::new();
        for j in i * len_per_commit..min((i + 1) * len_per_commit, x.len()) {
            tmp.push(x[j]);
        }
        commit_res.push(PedersenComScheme::commit(param, &tmp, r).unwrap());
    }
    commit_res
}

And here, we can see, for example, the Zero Knowledge circuit forargmax .

//used in mnist and cifar10 classification problem
#[derive(Debug, Clone)]
pub struct ArgmaxCircuit {
    pub input: Vec<u8>,
    pub argmax_res: usize,
}

impl ConstraintSynthesizer<Fq> for ArgmaxCircuit {
    fn generate_constraints(self, cs: ConstraintSystemRef<Fq>) -> Result<(), SynthesisError> {
        let _cir_number = cs.num_constraints();
        let argmax_fq: Fq = self.input[self.argmax_res].into();
        let argmax_var =
            FpVar::<Fq>::new_witness(r1cs_core::ns!(cs, "argmax var"), || Ok(argmax_fq)).unwrap();

        for i in 0..self.input.len() {
            let tmp_fq: Fq = self.input[i].into();
            let tmp = FpVar::<Fq>::new_witness(r1cs_core::ns!(cs, "argmax tmp var"), || Ok(tmp_fq))
                .unwrap();
            argmax_var
                .enforce_cmp(&tmp, Ordering::Greater, true)
                .unwrap();
        }

        println!(
            "Number of constraints for ArgmaxCircuitU8 Circuit {}, Accumulated constraints {}",
            cs.num_constraints() - _cir_number,
            cs.num_constraints()
        );
        Ok(())
    }
}

Finally, we can define a Private version of our initial network, adding the quantization and the ZKP circuits:

class Quant_Net(torch.nn.Module):
    def __init__(self):
        super(Quant_Net, self).__init__()
        self.l1 = nn.Linear(784, 128, bias=False)
        self.act = nn.ReLU()
        self.l2 = nn.Linear(128, 10, bias=False)
        self.quant = QuantStub()
        self.dequant = DeQuantStub()

    def forward(self, x):
        x = self.quant(x)
        x = self.l1(x)
        x = self.act(x)
        x = self.l2(x)
        x = self.dequant(x)
        return x

Results

We have experimented with the quantized Neural Network and got an accuracy of about 91.1% in the best of our running. There was a significant decrease in time when including the ZK circuits. For generating the Pedersen Commitments, the execution was 2,14785 seg. Generating the Proof using softmax ZK circuit was 126.54533 seg. Verification was faster than expected: 0.95123 seg.

Conclusions

We are reaching a point where we cannot ignore the privacy concerns that naturally arise in AI. We have a piece of very good machinery in Cryptography that can address some of that concerns.

Theoretically, the semantic closeness between Knowledge (as in Zero-Knowledge Proof) and Learning (as in Machine learning) is not a coincidence. An interesting relationship comes from the formalization of the learning concept in computer science.

We can go from Parity Learning to Learning with errors to get to some interesting Cryptography protocols whose assumption is hard, sometimes even for a quantum computer, and are based on the same Number Theory assumptions as the ones we discussed for ZKP.