Towards Fast and Scalable Private Inference

Two common linear computations in neural network inference are fully-connected and convolution layers, both of which can be expressed as matrix-vector multiplication. Here, the matrix represents the trainable parameters of the linear layer and the vector represents the input data. Given that linear transformations can be efficiently expressed as arithmetic circuits with bounded depth, a natural cryptographic primitive for privately computing the linear layers on encrypted inputs is HE.

HE is an encryption scheme that allows for computation directly on encrypted data without the need for decryption, thus preserving data confidentiality during the actual computation. Given two parties, a client (with some data) and a server (with a linear function), the client can first encrypt their data under a HE scheme and send their encrypted data to the server. The server can then homomorphically evaluate the linear function on the client’s ciphertext without ever decrypting the client’s data server-side. The resulting linear computation produces a ciphertext that can be sent back to the client who can decrypt the ciphertext locally.

HE introduces a large computational overhead of 4-6 orders of magnitude slowdown over their plaintext counterparts (1080 seconds for a single ResNet-18 inference (hybrid_PI, 10, 38)). Rather than directly using HE for computing linear layers, it is common to combine HE with Additive Secret Sharing (SS), another cryptographic building block that supports plaintext-level speeds for computing linear layers. In this HE+SS setting, linear layers are processed in two steps: a pre-processing or offline phase independent of the client’s input and an online phase dependent on the client’s input. In the pre-processing phase, the server performs a homomorphic evaluation of the linear layer on a randomly sampled and encrypted input to generate the client’s additive secret share. This pre-processing step enables the server to perform a simple secret share evaluation of the linear layer on the client’s actual data during the online phase.

Nonlinear layers are a crucial component in deep learning models, which include either an activation function or a pooling function. The activation function is typically chosen from several types of nonlinear functions, such as the widely-used rectified linear unit (ReLU) function. In convolutional neural networks, max-pooling functions are also frequently employed.

Generally, HE and additive SS enable private arithmetic computations, i.e., additions and multiplications over integer values. Garbled Circuits (Yao1986GC, 44) enable two parties to compute a Boolean function on their private inputs without revealing their inputs to each other. Unlike SS and HE, operating over Boolean gates enables the parties to jointly compute arbitrary function, making GCs a solution to privately compute nonlinear layers. We note that binary constructions of HE exist (e.g., TFHE (2020TFHE, 4)). However, these incur extremely high performance overheads with prior work. Evaluating single Boolean gate with HE can take 75-600ms to process (Stochastic_TFHE, 16, 28).

In GCs, a function is represented as a Boolean circuit. One party (the garbler) assigns random labels representing {0, 1} to each input wire of each gate, generates an encrypted truth table that maps the output labels to the gate’s input labels. The evaluator uses Oblivious Transfer (ObliviousTransfer, 33) to obtain labels corresponding to its inputs without revealing the values to the garbler, and then executes the circuit gate-by-gate. Finally, the evaluator shares the output labels with the garbler who maps them to plaintext values. Recent algorithmic optimizations are used to construct high-performance GCs (FreeXOR, 24, 45). More details can be found in (Half_Gate_rekey, 15, 31, 43).

Starting point. We briefly describe DELPHI (delphi, 29), a baseline hybrid PI protocol that uses HE and SS for linear layers, and GCs for ReLU layers. In the 2-party-computation setting, a client uses her data for inference on the server’s proprietary model, without learning the server’s model parameters. The DELPHI protocol consists of an offline phase, which only depends on the network architecture and parameters, and an online phase, which is performed after the client’s input is available. In this way, many expensive computations are moved to offline to improve online inference latency.

Offline Phase. First, the client sends their public keys to the server. At the $i$th layer, client and server sample random vectors $r_{i}$ and $s_{i}$ respectively. Then the client sends the encrypted random vector $E(r_{i})$ to the server. The server uses its model parameter $w_{i}$ to compute $E(w_{i}\cdot r_{i}-s_{i})$ homomorphically and returns to the client, thus the client holds the share $\langle y_{i}\rangle_{c}=w_{i}\cdot r_{i}-s_{i}$. The server also constructs GCs for each nonlinear ReLU operation and sent them to the client. The client obtains labels corresponding to its $\langle y_{i}\rangle_{c}$ and $r_{i+1}$ using OT, where $r_{i+1}$ is prepared for the next $(i+1)$th layer.

Online Phase. After the client’s input $x_{1}$ is available, the client sends the subtraction $x_{1}-r_{1}$ to the server. At the beginning of the $i$th layer, the server uses the subtraction to calculate its share of the layer output $\langle y_{i}\rangle_{s}=w_{i}(x_{i}-r_{i})+s_{i}$. To evaluate the ReLU layer, the server (garbler) sends the labels corresponding to its $\langle y_{i}\rangle_{s}$ to the client. The client uses its labels of $\langle y_{i}\rangle_{c}$ and $r_{i+1}$, together with the label of $\langle y_{i}\rangle_{s}$ to evaluate the garbled circuits ReLU$(\langle y_{i}\rangle_{c}+\langle y_{i}\rangle_{s})-r_{i+1}$. The GCs’ output labels will be sent to the server as they represent $(x_{i+1}-r_{i+1})$. At this point, the client holds the share $r_{i+1}$ and the server holds $(x_{i+1}-r_{i+1})$. They will similarly evaluate subsequent layers.

Addressing Storage: Storing the garbled ReLUs client-side during the offline phase inhibits the client from (at best) storing more than a few precomputes especially for deeper and larger networks. To overcome this limitation, the Client-Garbler protocol switches the role of the garbler and evaluator so that the server (now the evaluator) must store the ReLU GCs while holding the exact same security guarantees. This reduces the storage cost of the client device by $5\times$; rather than storing 41GB for TinyImageNet inference on ResNet-18, the client now only stores 8GB per inference. Outside of storage costs, Client-Garbler also reduces online phase latency as the powerful server is now evaluating the garbled ReLUs while the offline phase latency increases as the client must now perform garbling. This latency tradeoff is reflected in Figure 3 (+SysOpt).

Addressing Compute: Server-side HE evaluations of the linear layers during the offline phase account for a majority of the compute cost. As discussed in Section 2.3, the client and server iterate through each linear layer to generate secret shares that are used for fast, online phase linear evaluations. These secret shares are generated independently for each linear layer meaning each linear layer HE evaluation can be run in an embarrassingly parallel manner on the server which we call Layer-Parallel Homomorphic Evaluation (LPHE). Rather than performing each independent HE evaluation sequentially on a single core, each HE evaluation is allocated to a separate core server-side. For ResNet-18 on TinyImageNet, LPHE reduces the HE evaluation run-time from 1080s to just 141s. Across all datasets and networks, LPHE speeds up HE by $9.7\times$.

Addressing Communication: The transmission of garbled ReLUs from the garbler to the evaluator in the offline phase dominates the communication latency of the entire PI protocol. Furthermore, the massive storage penalty per ReLU (18.2KB) results in an asymmetry in the amount of data sent between the two parties. For example, in the Client-Garbler protocol, 83.5% of the total amount of data transferring is caused by uploading data to the server. Current 5G wireless standards allows for flexibility in the amount of bandwidth allocated to both uplink and downlink. By optimally allocating the bandwidth split (Wireless Slot Allocation), the communication latency of PI protocols can be reduces by 35%.

Putting these aforementioned system-level optimizations helps us not only reduces the storage cost of the client but also reduce both the offline and online phase latency. Figure 3 (+SysOpt) shows the latency breakdown with our optimized protocol. Compared to Figure 3 (Baseline), these optimizations reduce the end-to-end latency from 2050 seconds to 1052 seconds. Even after these optimizations, high compute costs from both the HE and GCs portions of the protocol dominate the runtime and necessitate custom architecture to further reduce latency.

Given the system-level optimization mentioned above, the GCs storage stress of the client side is largely relieved, but it raises the pressure on GCs computation. The GCs accelerator should be not only fast but also lightweight and energy efficient. It should be able to process large amounts of data and maintain high throughput. Here, we present HAAC (HAAC, 30), a novel hardware-software co-design that includes a compiler and hardware accelerator that combine to improve GCs performance and efficiency, making PI more practical.

Employing hardware-software co-design in GCs is appropriate, as the programs including all dependence, memory accesses, and control flow, are already known and fixed at compile time. Potential benefits in hardware can be realized through optimization on the software side. HAAC’s design philosophy is to keep hardware simple and efficient, maximizing area utilization of custom execution units and other circuits essential for high performance. Our approach leverages the properties of GCs to express arbitrary programs as streams, which simplifies hardware and enables complete memory-compute decoupling. The evaluation results demonstrate the effectiveness of the proposed approach in accelerating GCs and mitigating performance overheads.

HAAC hardware comprises Gate Engines (GE) execution units, see Figure 2. GE is deep, in-order pipeline that provide high performance potential, which ideally computes a gate per cycle. However, exploiting GEs parallelism while keeping hardware efficient is a challenge. This presents a prime opportunity for hardware-software co-design. HAAC’s compiler leverages Instruction-level Parallelism (ILP) to improve intra/inter-GE parallel processing, with instructions and constants streamed to each GE using queues. The compiler analyzes the leveled data dependence graph for the entire baseline program to expose all available ILP. This strategy, called full-reordering, works well to resolve data hazards, and significantly increases parallel GEs performance, but it can reduce the on-chip data reuse and burden the off-chip traffic. To better balance data reuse and ILP, HAAC also purposes another scheme called segment-reordering. Rather than computing the ILP graph for the entire program, we partition the program into contiguous parts (segments) and reorder instructions within each segment. This provides more instruction parallelism than baseline programs and generally captures more data reuse than full-reordering.

With the help of HAAC’s distinct on-chip memory subsystem, the high parallelism in GEs can be actually achieved instead of being blocked by input/output data latency. The gate inputs and outputs, called wires data, are more difficult as they do not follow a pattern. HAAC uses a scratchpad with multiple memory banks to capture wires reuse by tracking program execution, eliminating the need for costly hardware-managed caches and tagging logic. The on-chip scratchpad memory, called sliding wire window (SWW), stores a contiguous address range of wires, and the range increases as the program executes. We also leverage GCs input/output patterns and strides wires across SWW banks to reduce conflicts. Renaming is a complementary compiler pass that sequentializes gate output wire addresses according to program order.

The SWW and renaming combine to filter off-chip accesses, as recently generated wires are often soon reused when they are still on-chip, with the performance benefits of a cache and the efficiency of a scratchpad. HAAC’s co-design also enables complete gate execution and off-chip accesses decoupling. Used wires, marked by the compiler, will not go off-chip by Eliminating Spent Wires (ESW), which elides redundant writes to off-chip memory. The unused wires, or Out-of-range Wire (OoRW), will be stored in the off-chip memory. The compiler knows which wires will be OoR and can push them on-chip to an OoRW queue, and the GEs will to check the OoRW queue if the inputs are not on-chip. The optimizations allow the overlap of computation and off-chip accesses, hiding the latency of data movement.

We evaluate HAAC thoroughly with benchmarks from VIP-Bench (VIP-Bench, 2). With 16 GEs, a 2MB SWW on-chip memory, and DDR4, HAAC provides an average speedup of 589$\times$ than an Intel Core i7-10700K CPU (2627$\times$ with HBM2) in 4.3mm². Figure 3 (+HAAC) shows that after accelerating GCs with HAAC, we further reduce the total latency for a private inference on ResNet-18 by 39%. The percentage of GCs in the computation latency is already very low, leaving the HE evaluation as a large occupation.

We now turn to the final optimization: accelerating expensive HE operations required for linear layers in the offline phase. Like RELUs, simple plaintext operations see large overheads in the encrypted domain. Specifically, linear transformations map to a series of complex ciphertext rotations and key switching operations, which dominate the runtime of HE, requiring frequent applications of the (inverse) number theoretic transform, or (i)NTT. Fortunately, this implies high degrees of data-level parallelism, which can be efficiently exploited with the right architectural design.

With this in mind, we now discuss our Ring Processing Unit (RPU) (RPU, 40). The RPU implements B512, a vector ISA that has been specifically designed to cater to common HE operations while also remaining well-suited to general-purpose programming. The decision to develop a vector architecture and ISA, rather than fixed ASIC hardware, gives us the reprogrammability needed to remain viable as HE algorithms continue to develop and mature.

The vector architecture state includes 64 128-bit vector registers, 64 128-bit scalar registers, a 4 MiB vector data memory (VDM) expandable up to 32 MiB, and a 2 MiB scalar data memory (SDM). Additionally, there is an Address Register File (ARF) for indirect memory access, a Modulus Register File (MRF), and several control registers for increased flexibility. We operate on fixed vector lengths of 512 elements to balance microarchitectural concerns with the power-of-two input sizes common to both (i)NTT and HE kernels. Instructions are grouped into four categories: compute, memory, shuffle, and control. Compute instructions perform point-wise modular computations between two vectors or a vector and a scalar value. Memory operations interact with the VDM and SDM to bring data to and from the vector and scalar register files respectively. We support four addressing modes in memory instructions to efficiently handle the complex access patterns in the (i)NTT and other HE workloads. Shuffle instructions are heavily used by (i)NTT and are efficiently supported in the microarchitecture. Control instructions act on the Control Register File (CRF) and drastically reduce our required code size. Overall, B512 introduces 28 instructions to meet the needs of HE-specific applications without sacrificing the generality that traditional vector operations provide.

The RPU is designed for general ring processing with high performance by taking advantage of regularity and data parallelism. We achieve this balance by designing explicitly managed hardware to elide the high costs and complexity of caches, dynamic scheduling logic, and prediction, and task the compiler with handling scheduling and data movement at compile time. Figure 2 presents the RPU microarchitecture, which consists of a front-end to handle instruction fetching, decoding, and control logic, and a backend, which provides the high-performance hardware needed to efficiently perform our HE-tailored vector operations. The front-end includes three decoupled queues that operate independently on compute, memory, and shuffle instructions. Once an instruction is in its respective queue, it can run in parallel with other instruction types, and the microarchitecture guarantees that data hazards are avoided. The parallel execution via these decoupled pipelines is key to achieving high performance with general-purpose processing as it masks much of the data movement time.

Compute instructions are passed to computational units that we denote as high-performance large arithmetic word (LAW) engines (HPLEs). HPLEs each operate on a slice of the 512 element vector whose size is determined by the number of parallel HPLEs (i.e., lanes) in our vector architecture. Each HPLE consists of a 128-bit modular multiplier, adder, subtractor, and two comparator units. HPLEs can be modified to support bit-serial computation, trading off performance for area.

We use the SPIRAL (Spiral, 8) to map instructions onto the RPU and automatically generate high-performance B512 programs. SPIRAL has a rich library of transformations and optimizations that can expertly generate high-performance code across various platforms and kernels, especially in linear transforms such as the (i)NTT. We evaluate and characterize RPU using a detailed cycle-level simulator that is parameterized to consider a range of IP, namely modular multiplier designs, number of HPLEs, and VDM partitioning strategies. This enables rapid design space exploration to quantify design decisions. The simulator is further verified with a complete RTL implementation of the RPU. We show that our most efficient parameter choices can accelerate varying size (i)NTTs by up to $1485\times$ over a traditional 32-core 2.5GHz AMD EPYC 7502 CPU.

Figure 3 (+RPU) shows that the HE evaluation latency is largely eliminated. Overall, by performing the system-level optimization, together with the hardware accelerators HAAC and RPU, we significantly reduced the end-to-end private inference latency by 76% against the baseline. The computation overhead is inconspicuous compared to the communication, which will be resolved by developing faster bandwidth and network technology in the future.