It becomes clearer coding with AI is increasing productivity significantly, if done right. Some tasks are heuristic in nature, like generating images or summarizing articles. Concretely, the success criterion is more similar to "it looks good" and "it's useful" than anything else. In coding you can often do better. For example, you can have tests which guarantee flows of the program act in a certain way. A function receiving certain inputs produces expected outputs. When operating on pure functions, those that don't interact with or affect other parts of a bigger system, the test is robust. You can make changes and optimize a pure function without worrying too much about whether you broke the system. In general software, these tests provide good value in their local vicinity. When looking at the bigger system that interacts with many moving parts, such as databases, network, storage and more, you want global guarantees, which you can achieve through methods like integration tests that span entire systems.
What about cryptography? In cryptography the APIs are often stateless and simple, even though the math can be hard. For example, signatures have something like:
- $\mathsf{GenerateKeys}() \rightarrow (sk, pk)$: generate a secret key and public key.
- $\mathsf{Sign}(sk, msg) \rightarrow \sigma$: sign a message.
- $\mathsf{Verify}(pk, msg, \sigma)$: verify a signature on a message.
This looks the same for ECDSA, Schnorr, BLS and many other schemes.
Optimizing polynomial commitment schemes
Let's talk about modern proof systems that are based on polynomial commitment schemes, or PCSes. In a PCS you have two main phases:
- $\mathsf{Commit}(p) \rightarrow C$: commit to a polynomial.
- $\mathsf{Evaluate}(C, x) \rightarrow (z, \pi)$: evaluate the polynomial committed in $C$ at point $x$ and return the evaluation $z$ with a proof $\pi$.
In many modern proof systems, the polynomial commitments are taking the bulk of the work, and so it's important to optimize them. This is, for example, what we (Andrija, Ron and I) set out to do when working on Bolt, exploring a point in the tradeoff space where you aim to get the fastest commitment at the expense of a larger, though practical, proof size.
One question then comes up - can you use AI safely in this situation to optimize? The answer is yes! With caveats.
Let's look at the commitment phase of schemes based on error-correcting codes. The general pattern is that you encode the polynomial, commit to the encoding and then query the committed encoding a sufficient amount of times to show the commitment corresponds to a unique polynomial (or a small set of polynomials). The security of this process mainly depends on the distance of the code, which is a property you analyze outside the implementation. It involves randomness in choosing the query points, which is usually achieved using Fiat-Shamir for non-interactivity.
Given this, a nice test would be to generate one good encoding, save the resulting commitment and start optimizing. As long as the commitment stays the same, you're good! You know that you're encoding correctly. The danger with this approach is that if you give this task to an AI and you're not careful, it might "overfit" on this specific task and reproduce the commitment in invalid and fast ways. For example, by precomputing it once and saving it in the file.
So you can do even better. You can have a slow and obviously correct implementation that dynamically generates test cases for you, which increases your confidence that an optimized version is correctly encoding each time.
Sounds good in theory... Is it practical?
In Bolt, this was the case. There was an initial Julia implementation, based on Andrija's original work on Ligerito, which he then extended to Bolt's commitment.
First, we took that one and used AI to first convert it to Rust, generating two-way equivalence tests along the way for all the components we could, e.g. finite field implementation and FFT. By two-way equivalence I mean the following: Julia generates test cases for these components that Rust needs to reproduce correctly, and Rust generates tests cases that Julia needs to reproduce. This ended up working, including the Bolt commitment itself. This requires manual verification, at least on some level, as the tests can make wrong assertions. This happened to me in the past, where tests would not even have an assert_eq.
As a side note, the Bolt commitment is a combination of LDPC and RS codes. LDPC codes end up essentially being a bunch of linear combinations of randomly sampled vectors from the data matrix. RS codes are well known.
Having that strong base, we continued by starting to throw as many ideas as we had at the AI, while making sure that the commitment stays the same. This worked extremely well. We managed to try many more ideas than we would have if we had to do it ourselves, and we had the confidence that the implementation is correct since the commitment produced the same result as the slow implementation. This didn't always work, since our hunches weren't always correct, and the AI hunches' weren't always correct either. In a few instances, the AI would say something like "yes, this is going to be amazing" which ended up having little effect on performance.
The most successful strategies ended up being the following, where our target machine was a Mac with an M processor:
-
Cache optimizations - smartly load soon-to-be-needed elements into the L1 cache, which is small, and even more so preload future elements into the bigger L2 cache so that loading into L1 cache will be faster.
-
Use the SIMD vector instructions, which works well since we're adding, and multiplying by scalar, whole columns.
-
GPU acceleration using Metal for encoding and hashing, which we had no material experience with.
-
Assembly optimizations, using carryless multiplication and other instructions.
-
Hardware-accelerated instructions for hashing. For example, for SHA256, it's done through using the ring library. Underneath, there are instructions such as
SHA256Hthat implement parts of SHA256 hashing. -
Pipelining the work between encoding and committing, since Bolt's code is systematic, meaning that a big chunk of the encoding is hashing the message itself.
This resulted in more than an order of magnitude improvement in performance over the Julia implementation, through MANY rounds of prompting and reviewing.
The human-in-the-loop is still important. For example, one optimization that the AI came up with was great, though I didn't review it well and it tripped me. It was around how to embed an 8-bit field element into a 32-bit field element with a specific representation. Without getting too much into the details, the right way to do it and still preserve the distance of our code would be to use an isomorphism between the representation of the 32-bit field as a vector of four 8-bit elements to the desired representation. What the AI suggested, not knowing that it hurts the distance of our code, is reinterpreting the bits of 8-bit element as bits of the 32-bit element. This is some sort of embedding, though it basically makes the element be as any other 32-bit element, rather than having the unique property of coming from an 8-bit element.
That sounds awesome, will this always work?
Unfortunately no! This kind of commitments is especially suited for optimizations. Other kinds of systems may not be. For example, ZK circuits can be wrong in subtle ways. You have the goal of reducing constraints to optimize the prover's performance, though a missing constraint in a circuit still yields the same functional result, though can be completely broken. Even a single missing constraint can cause complete break, as can be seen over and over in audits and real exploits.
Another instance where the strategy is successful is optimizing a cryptographic protocol while preserving soundness, if done carefully, since this is more subtle. This works if you have a slow and correct implementation that produces the same elements, bit by bit. In interactive public-coin protocols, that are then transformed into non-interactive ones using Fiat-Shamir, this takes the form of reproducing the transcript - the set of messages the prover and verifier sent each other during the protocol. If these remain the same, and your abstraction is tight enough to guarantee there are no other side channels, you can be confident the optimizations were a net improvement.
You can generalize this notion even further. Whenever you can define a verifier that guarantees soundness, and it still works while you optimize, you're in a safe situation. Note that other properties such as completeness or zero-knowledge can still be affected, unless you also have a way to model their verification as well. This is especially touchy, since optimized code could be incomplete. For benchmarks it's fine, and for production code it's not.
This risk of a program not acting according to a specification can be mitigated, with non-trivial effort, by methods like formal verification. Even more concretely, Gregor and the team have been working for the last couple of years on Clean, which brings formal verification to ZK circuits. We're using it in a few different projects now.
Now, let's say you have a circuit for SHA256 defined in Clean, can you then go back to focus on optimizing it without worrying you broke soundness? Yes! That's exactly what Clean gives you. No matter how you optimize, the specification will guarantee soundness is preserved.
Because it's such a powerful mechanism, Giorgio and Mathias have launched a competition to optimize ZK circuits that are specified in Clean. You want to do it by hand? Go ahead. A legion of AIs? Also great. Just make sure it still adheres to the specification and you're good :)
Check it out here and see how far you can get. There have been pretty big optimizations contributed to the circuits already - both for prime fields, around hashes and scalar multiplication, and binary fields, around hashes.
A couple of other cool efforts include snark.fast and lighter.fast that work with the same idea of optimize the prover while keeping the verifier the same.
Is there something general to take away?
The main piece of advice I have is to focus on what are the right anchors to make sure your AI solves the right problem. For a commitment scheme, it can be the commitment hash. For a ZK circuit, it can be the verifier or the formal specification. For an AI model, it could be a loss metric. You can be creative with these, just make sure that the anchor is effective. We've talked a bit about this in a recent episode of ZKPodcast with Benedikt Bünz, where Benedikt, Ron and William used AI heavily to optimize the implementation of their new paper Flock, which is a proof systems for extremely fast proving for bit-twiddling hashes. We also discussed AI-assisted work more generally.
Acknowledgements
Thanks to Mathias for a comment about the guarantees a verifier gives you. Thanks to Stefanos for fixes and comments about the reliability of equivalence tests and hardware acceleration. Thanks to Ron and Andrija for reading through the post.