DeepReShape: Redesigning Neural Networks for Efficient Private Inference

TL;DR

DeepReShape introduces ReLU-equalization to reorder ReLUs by their accuracy criticality and designs HybReNets that balance ReLUs and FLOPs under private inference constraints. The framework combines ReLU-equalization with ReLU-reuse and coarse-grained pruning to achieve large FLOP reductions while maintaining or improving accuracy, translating to substantial end-to-end latency savings in HE+GC private inference. The results demonstrate state-of-the-art Pareto performance on CIFAR-100 and TinyImageNet, with significant runtime improvements even when FLOPs are drastically reduced, and reveal the importance of network selection and stage-wise channel distribution for PI efficiency. The work offers a general design principle for PI-efficient networks, with potential applicability beyond private inference to other privacy-preserving ML applications and nonlinearities.

Abstract

Prior work on Private Inference (PI) -- inferences performed directly on encrypted input -- has focused on minimizing a network's ReLUs, which have been assumed to dominate PI latency rather than FLOPs. Recent work has shown that FLOPs for PI can no longer be ignored and incur high latency penalties. In this paper, we develop DeepReShape, a technique that optimizes neural network architectures under PI's constraints, optimizing for both ReLUs and FLOPs for the first time. The key insight is strategically allocating channels to position the network's ReLUs in order of their criticality to network accuracy, simultaneously optimizes ReLU and FLOPs efficiency. DeepReShape automates network development with an efficient process, and we call generated networks HybReNets. We evaluate DeepReShape using standard PI benchmarks and demonstrate a 2.1% accuracy gain with a 5.2$\times$ runtime improvement at iso-ReLU on CIFAR-100 and an 8.7$\times$ runtime improvement at iso-accuracy on TinyImageNet. Furthermore, we investigate the significance of network selection in prior ReLU optimizations and shed light on the key network attributes for superior PI performance.

Figures (23)

• Figure 1: HybReNet outperforms state-of-the-art (SOTA) ReLU-optimization methods SENetskundu2023learning, SNLcho2022selective, and DeepReDucejha2021deepreduce, achieving higher accuracy (CIFAR-100) and significant reduction in FLOPs while using fewer ReLUs (Table \ref{['tab:ParetoResNet18C100']} illustrates the Pareto points specifics).
• Figure 2: Depiction of architectural hyperparameters and feature dimensions in a four stage network. For ResNet18 $m$ = 64, $\phi_1$=$\phi_2$=$\phi_3$ =$\phi_4$=2, and $\alpha$=$\beta$=$\gamma$=2.
• Figure 3: (a) Homogeneous channel scaling in StageCh networks enables superior ReLU efficiency compared to uniform channel scaling in BaseCh networks; however, (b) the accuracy in StageCh networks tends to plateau unpredictably. (c,d) Unlike uniform channel scaling, homogeneous scaling reduces the proportion of least-critical ReLUs in StageCh networks. (e,f) Each network stage affects ReLU and FLOPs efficiency differently, requiring heterogeneous channel scaling for optimizing both ReLUs and FLOPs for efficient PI.
• Figure 4: (a) Unlike StageCh networks, once the network's ReLUs are aligned in their criticality order, here at point ($\alpha$, $\beta$, $\gamma$)=(5, 5, 3), increasing $\alpha$ does not alter their relative distribution. (b,c) ReLUs' criticality-aware network widening method saves 2$\times$ FLOPs by regulating the FLOPs in deeper layers while maintaining ReLU efficiency over a wide range of ReLU counts.
• Figure 5: ReLU optimization, whether coarse or fine-grained, performance exhibits significant disparities based on the input networks.