> ## Text from Telegram
1. Detailed spec of
1.1 ONNX standard-compliant circuits (as listed in here https://github.com/onnx/onnx-mlir/blob/main/docs/SupportedONNXOps-cpu.md) formal description. Just like we did once in here: https://github.com/NilFoundation/solana-state-proof/blob/master/docs/design/main.pdf. Some of those from ONNX standard are already present on our side - let's coordinate about this.
1.2 Analysis about if there is sense to reduce from MLIR to LLVM IR.
1.3 If it makes sense, then how to lower MLIR to LLVM IR efficiently enough and which extensions LLVM IR has to have.
1.4 More efficient circuits definition done for a particular example use case. (Example should be big and ambitious one for this to worth it to you guys, so let's make it some trivial LLM-alike one? We have some circuit estimations about such one)
1.5 Analysis about necessary proof system efficiency improvements.
## Questions
* Cost metrics for their proof system, how does adding rows/columns affect size/time? [R]
* List of Operations in MLIR (e.g, arith.cmpf, affine.for, affine.max) [R]
* Is pattern matching even feasible at the MLIR level?
* Can the ONNX to MLIR compiler be modified easily?
* MLIR has the concept of hardware accellerators, can we somehow "hook" in there?
* Existing literature for zkML in plonkish arithmetization systems [R]
* https://arxiv.org/pdf/2210.08674.pdf
* https://drive.google.com/file/d/1tylpowpaqcOhKQtYolPlqvx6R2Gv4IzE/view
* https://hackmd.io/Y7Y79_MtSoKdHNAEfZRXUg
* Approximation of exp, etc (see Wiki) [R]
* Approximation in training -> Cannot use existing models as easily.
* Look into quantization more
* Fixed-point planning [R]
* Select fixed point scale factor
* Algorithms for Add, Mul, BitExtract, Rescale
* Floating-point investigation [R]
* https://www.plonk.cafe/t/draft-plonk-intermediate-opcode-representation/45
* Other methods for floating points in (sign, mantissa, exponent) format
* https://eprint.iacr.org/2021/1063.pdf (Cairo, sec 8.3, p45)
* Specialized Functions [R]
* ReLU (Max, signextract, etc)
* Approximated ReLu with Poly already in training <https://arxiv.org/pdf/2011.05530.pdf>
* Math (exp, sin, cos, tanh, ...)
* Approximate via rational methods <https://xn--2-umb.com/22/approximation/>
* <https://xn--2-umb.com/22/exp-ln/>
* Model constants
* Where to put them in the trace?
* Generic stuff, like how big is a trace with just the model constants in there? is this already a problem, re Cost Metrics above.
* GPT-2 as a case study:
* <http://jalammar.github.io/illustrated-gpt2/>
* Token embedding is 50207x768 = Already 38M...
* Is needed for each input token + at the end of each of the 12 layers
* Matrix needs to be repeated in the trace probably...
* Also this is only for 1 token...gpt2 can do up to 1000
* Alternative LLMs: <https://huggingface.co/>
* BERT?
* Mini-Bert: https://huggingface.co/prajjwal1/bert-mini/tree/main
* MiniLM <https://huggingface.co/microsoft/MiniLM-L12-H384-uncased>, <https://huggingface.co/nreimers/MiniLM-L6-H384-uncased>
* Recursion, Folding, etc...
* Investigate Methods used in RiscZero
* <https://geometry.xyz/notebook/sangria-a-folding-scheme-for-plonk>
* Cannot do lookup arguments atm
* ONNX layer
* Existing libraries to parse this (in C++?)
* Selection of important blocks from <https://github.com/onnx/onnx-mlir/blob/main/docs/SupportedONNXOps-cpu.md>
* Existing libraries to work with ONNX (the python lib can do size inference for example, which is pretty much required for us to work with the model later on, but this can be "baked into" the ONNX model)
* Assigner Formats
* Reverse Engineer their format, ask for documentation
* Enables us to write assigner in other language
* Need to keep in sync with C++ version
* Basic Matrix Multiplication
* Investigate layouting matrix in trace
* Advanced
* Existing zkML Code stuff
* <https://github.com/0xZKML/zk-mnist> MNIST in Circom
* <https://github.com/socathie/circomlib-ml> Basic ML building blocks in Circom
* https://github.com/worldcoin/proto-neural-zkp
* <https://github.com/ddkang/zkml> Probably pretty advanced? 8s for MNIST proving, blog post below is even targeting gpt...
* https://medium.com/@danieldkang/verified-execution-of-gpt-bert-clip-and-more-6acb693fd55f
* https://arxiv.org/pdf/2210.08674.pdf
* https://github.com/lyronctk/zator
# Spec:
# Introduction
The goal of this zkML project is to extend the [zkLLVM project](https://github.com/NilFoundation/zkllvm) with capabilities to prove machine learning inference.
# Milestones
1. Detailed Spec. Report (2 months, 3 developers)
2. ONNX-frontend implementation (TODO)
3. Example use case implementation (TODO)
## Milestone 1: Detailed Spec. Report
In the first phase of the project we will investigate the following questions:
- The base will be ONNX models:
- We will investigate the efficient arithmetization of the most important/used operations of the [ONNX standard](https://github.com/onnx/onnx-mlir/blob/main/docs/SupportedONNXOps-cpu.md)
- We will base our selection of supported operations on an investigation of the operations used in the most-used ONNX models on huggingface.co
- We will investigate which custom gates we would need to efficiently arithmetize these standard operations
- e.g., dot-product components that can be re-used for matrix multiplication and convolutions
- These components are built with respect the the floating/fixpoint arithmetization findings below
- We will provide detailed trace layouts and constraints for the components to quantify their concrete costs
- Floating/Fixpoint Arithmetization:
- Most state-of-the-art literature suggests using fixed-point representations of the involved floating point values. We will investigate how to efficiently represent them in the context of ML models and how to arithmetize rescaling operations.
- We will investigate ways to combine the fixed-point representation with the custom gates above, e.g., rescaling could happen after the dot-product component instead of directly rescaling each multiplication result.
- We will investigate quantization processes for existing models or potential support for already quantized models.
- Compilation chain:
- We will investigate the best place to make our customiztions:
- ONNX graph -> ONNX-MLIR dialect:
- ONNX MLIR is a representation of the ONNX graph in MLIR, with several optimizations already applied by the ONNX-MLIR backend. This is a high-level (op-level) view of the ONNX graph that would allow us to build more efficient components.
- Allows for the highest degree of freedom in arithmetizing the blocks
- ONNX graph -> ONNX-MLIR -> Affine MLIR dialect
- The affine MLIR dialect is a breakdown of the ONNX MLIR into krnl-dialect which is then mapped to affine-dialect
- Affine dialect works with low-level loops and floating point math operations, which is harder to map to specific, more efficient custom gates
- ONNX-MLIR hardware accelerators:
- ONNX-MLIR has the concept of hardware accelerators for specific functions.
- This could also be a an alternative point to hook into the existing ecosystem
- We will conduct experiments to justify our findings in the report
- Showcase on ML model:
- We will showcase our findings by applying them to a specific ML model.
- Could be GPT-2, MINI-BERT, or similar
- These models already have millions of parameters which need to be efficienlty included in the execution trace
- We will investigate the tradeoffs in efficient circuit layouts (i.e., optimizing the prover cost metric)
- We will investigate useful extensions of the proof system that would enable more efficient implementation of specific ONNX operations
- We will investigate the feasibility of using smaller field sizes in the proof system for increased efficiency
- We will investigate if recursion can be applied to improve efficiency
## Milestones 2: Finishing ONNX Frontend with supported ONNX operations
Phase 2 will start with a retrospective meeting on the findings of the report in phase 1 and how they influence and change the plans layed out below for phase 2 and 3.
After feedback on the findings of the report and integration of potential changes into the work packages below, we will work on implementing our findings into an ONNX frontend.
The frontend will consist of the following parts:
- A CLI application that can accept ONNX models and inputs to these models and
1. Compiles them to ONNX-MLIR dialect identified in the report.
2. Evaluates the model on the input in the chosen (e.g., fix-point) representation
3. Performs the trace layouting for the model circuit
4. Generates the assignment for the trace
5. Hands the trace and assignment off to the low-level prover to create the proof
6. Can verify produced proofs locally
To facilicate the above frontend we will provide:
- Implementation of a parser for the ONNX-MLIR dialect identified in the report
- This parser will essentially be a counterpart to the assigner, focusing on layouting and assigning the circuit components for the ONNX operations
- The parser and components will follow findings of the report with respect to the fixed-point representation
- Implementation of the custom components identified in the report in the low-level blueprint library
- The individual components will follow the findings of the report
- These components will allow the proving of all ONNX operations currently supported in ONNX-MLIR, taking into account limitations of the ONNX-MLIR layer and any limitations identified in the report
- https://github.com/onnx/onnx-mlir/blob/5924af533182902d568d4765bc26be1f758bb1d5/docs/SupportedONNXOps-cpu.md
- Unit Test for the individual ONNX operations
- For each implemented ONNX operator at the MLIR level, we will provide one or more unit test ensuring the functionality of the individual operations
- Benchmarks and concrete component sizes for different operators
- Examples
- We will provide a few end-to-end examples showcasing the functionality of the system for a few chosen well-known models.
## Milestone 3: Example use case implementation
This milestone focuses on the implementation of the use case considering findings in the report from Milestone 1, using the functionality implemented in Phase 2.
- Implementation of a concrete use-case to showcase the functionality of the ONNX frontend
- The concrete use case is planned to follow the analysis of the large model from the report in Phase 1 (GPT-2, MINI-BERT, etc.)
- However, potential changes due to the findings of the report might lead to a different choice of use-case
- The implementation will consist of a demonstrator using the functionality implemented in Phase 2 to prove inference of a large machine learning model.
- This demonstrator will consist of a prover and verifier and documentation on how to use them (and what kind of resource requirements exist)
- Deployment via a Docker image will be supported
- We also aim to provide an already computed proof so more resource-limited clients can just execute the verifier