A system that allows secure machine learning inference in a two-party setting. The system automatically searches for an optimal ternary neural network architecture, trains it, and then builds the requisite files for secure inference.
The algorithms are explained in the following research paper: Aggarwal, Anshul, Trevor E. Carlson, Reza Shokri, and Shruti Tople. "SOTERIA: In Search of Efficient Neural Networks for Private Inference." arXiv preprint arXiv:2007.12934 (2020).
Make sure you have docker installed before you set up the docker image. Navigate to the code directory containing the Dockerfile, and execute:
$ docker build ./
This will build the Docker image, and install all the requisite tools required to execute the code. This needs to be run only once.
It may take a few minutes to complete. An image ID will be returned. Create a container and run bash over the image, using the command:
$ docker run -it --runtime=nvidia <image-id> /bin/bash
Replace <image-id> with the image ID obtained when the docker image is built.
The first step is the DARTS Architecture Search, to find a suitable architecture for a given dataset, and then train it. Please see the code and README file in Soteria/Architecture_Search for more details.
Included is code for training any manually defined TNN on MNIST and CIFAR10 datasets, without going through the Architecture Search process. Please see the code and README file in Soteria/TNN for more details.
We try to reproduce the results for existing work XONN. Please refer to the directory Soteria/TNN/XONN/ for more details.
To build the model in Verilog and execute the secure inference protocol, please refer to the code and README file in Soteria/Verilog_constructor for more details. You will also find some automated scripts to run the tests and generate the outputs in the directory. Further, scripts used to generate the different results in the paper are also provided, and references to those scripts are made in the README file in the directory.
Code from the following sources was used and modified to build parts of our system.
GitHub: https://github.com/MandyMo/DARTS
Paper: Hanxiao Liu, Karen Simonyan and Yiming Yang. DARTS: Differentiable Architecture Search, ICLR 2019.
GitHub: https://github.com/itayhubara/BinaryNet.pytorch
Paper: Itay Hubara, Matthieu Courbariaux, Daniel Soudry, Ran El-Yaniv and Yoshua Bengio. Binarized Neural Networks, NIPS 2016.