[Paper Review] Efficient ML Models for Practical Secure Inference

Vinod Ganesan, Anwesh Bhattacharya+4 more|arXiv (Cornell University)|Aug 26, 2022

Radiomics and Machine Learning in Medical Imaging4 citations

TL;DR

This paper proposes a framework to make secure inference practical for deep learning models by optimizing model architecture and computation for secure 2-party computation. It introduces a crypto-friendly X-operator, efficient convolution factorization, and Winograd algorithm integration, reducing communication costs by up to 30× with minimal accuracy drop (AUC loss ≤0.02) on DenseNet-121 for chest X-ray diagnosis.

ABSTRACT

ML-as-a-service continues to grow, and so does the need for very strong privacy guarantees. Secure inference has emerged as a potential solution, wherein cryptographic primitives allow inference without revealing users' inputs to a model provider or model's weights to a user. For instance, the model provider could be a diagnostics company that has trained a state-of-the-art DenseNet-121 model for interpreting a chest X-ray and the user could be a patient at a hospital. While secure inference is in principle feasible for this setting, there are no existing techniques that make it practical at scale. The CrypTFlow2 framework provides a potential solution with its ability to automatically and correctly translate clear-text inference to secure inference for arbitrary models. However, the resultant secure inference from CrypTFlow2 is impractically expensive: Almost 3TB of communication is required to interpret a single X-ray on DenseNet-121. In this paper, we address this outstanding challenge of inefficiency of secure inference with three contributions. First, we show that the primary bottlenecks in secure inference are large linear layers which can be optimized with the choice of network backbone and the use of operators developed for efficient clear-text inference. This finding and emphasis deviates from many recent works which focus on optimizing non-linear activation layers when performing secure inference of smaller networks. Second, based on analysis of a bottle-necked convolution layer, we design a X-operator which is a more efficient drop-in replacement. Third, we show that the fast Winograd convolution algorithm further improves efficiency of secure inference. In combination, these three optimizations prove to be highly effective for the problem of X-ray interpretation trained on the CheXpert dataset.

Motivation & Objective

• To address the impractical communication overhead in secure inference for large deep learning models.
• To identify and optimize the primary bottlenecks in secure inference, particularly large linear layers and convolution operations.
• To enable practical, privacy-preserving machine learning inference in high-stakes domains like healthcare using secure 2PC.
• To integrate efficient, secure-friendly operators and algorithms into the CrypTFlow2 framework for end-to-end secure inference compilation.
• To demonstrate that significant efficiency gains are achievable without sacrificing model accuracy or correctness.

Proposed method

• Designed and implemented the X-operator as a secure, efficient drop-in replacement for standard convolution layers, combining factorization, shuffle operations, and addition-heavy computation.
• Applied factorization to convert dense convolutions into depthwise-separable convolutions, reducing communication by 4.07× on average.
• Integrated the fast Winograd convolution algorithm into CrypTFlow2’s Athos compiler frontend to accelerate dense convolutions with no accuracy loss.
• Used shuffling operations to optimize channel-wise transformations in factorized convolutions, reducing communication by 1.1× on average.
• Combined the X-operator with Winograd and other optimizations to achieve cumulative efficiency gains across multiple model architectures.
• Evaluated all optimizations end-to-end in CrypTFlow2, ensuring bit-wise equivalence between clear-text and secure inference outputs.

Experimental results

Research questions

• RQ1 What are the primary communication bottlenecks in secure inference for large models like DenseNet-121?
• RQ2 Can efficient clear-text inference operators be adapted to reduce secure inference costs without compromising accuracy?
• RQ3 How effective is the X-operator in replacing standard convolutions for secure inference in terms of communication and latency?
• RQ4 To what extent can the Winograd algorithm reduce communication costs in secure inference for dense convolutions?
• RQ5 Can the combination of architectural choices, custom operators, and algorithmic optimizations achieve practical secure inference for medical image analysis?

Key findings

• The X-operator reduced communication cost by up to 30.8× (19.1× for ResNet-18, 30.15× for MobileNetV3) compared to baseline DenseNet-121, with only a 0.02 AUC drop.
• Factorized convolutions reduced communication by 4.07× on average, with ResNet-18 achieving 12.41× reduction (R′F vs DD).
• The Winograd algorithm reduced communication by 2.12× on average for dense convolutions, with no accuracy loss.
• Shuffle operations provided marginal gains, reducing communication by 1.1× on average over factorized convolutions.
• The combination of X-operator and Winograd reduced communication by 30.8× (SXW) with only 0.02 AUC drop compared to baseline DenseNet-121.
• All optimized models maintained bit-wise equivalence with clear-text inference, ensuring correctness and robustness.