[Paper Review] Accurate Low-Degree Polynomial Approximation of Non-polynomial Operators for Fast Private Inference in Homomorphic Encryption

Jianming Tong, Jingtian Dang+4 more|arXiv (Cornell University)|Apr 4, 2024

Cryptography and Data Security4 citations

TL;DR

This paper proposes SMART-PAF, a framework that enables accurate and efficient private inference in homomorphic encryption by replacing non-polynomial operators (e.g., ReLU, MaxPooling) with low-degree polynomial approximations (PAFs). It introduces four synergistic techniques—Coefficient Tuning, Progressive Approximation, Alternate Training, and Dynamic/Static Scaling—to stabilize training and reduce accuracy degradation, achieving up to 13.64× accuracy improvement and 14.9× speedup over prior methods on ResNet-18 with ImageNet-1k.

ABSTRACT

As machine learning (ML) permeates fields like healthcare, facial recognition, and blockchain, the need to protect sensitive data intensifies. Fully Homomorphic Encryption (FHE) allows inference on encrypted data, preserving the privacy of both data and the ML model. However, it slows down non-secure inference by up to five magnitudes, with a root cause of replacing non-polynomial operators (ReLU and MaxPooling) with high-degree Polynomial Approximated Function (PAF). We propose SmartPAF, a framework to replace non-polynomial operators with low-degree PAF and then recover the accuracy of PAF-approximated model through four techniques: (1) Coefficient Tuning (CT) -- adjust PAF coefficients based on the input distributions before training, (2) Progressive Approximation (PA) -- progressively replace one non-polynomial operator at a time followed by a fine-tuning, (3) Alternate Training (AT) -- alternate the training between PAFs and other linear operators in the decoupled manner, and (4) Dynamic Scale (DS) / Static Scale (SS) -- dynamically scale PAF input value within (-1, 1) in training, and fix the scale as the running max value in FHE deployment. The synergistic effect of CT, PA, AT, and DS/SS enables SmartPAF to enhance the accuracy of the various models approximated by PAFs with various low degrees under multiple datasets. For ResNet-18 under ImageNet-1k, the Pareto-frontier spotted by SmartPAF in latency-accuracy tradeoff space achieves 1.42x ~ 13.64x accuracy improvement and 6.79x ~ 14.9x speedup than prior works. Further, SmartPAF enables a 14-degree PAF (f1^2 g_1^2) to achieve 7.81x speedup compared to the 27-degree PAF obtained by minimax approximation with the same 69.4% post-replacement accuracy. Our code is available at https://github.com/EfficientFHE/SmartPAF.

Motivation & Objective

• Address the high latency and accuracy degradation in fully homomorphic encryption (FHE)-based private inference caused by non-polynomial operators like ReLU and MaxPooling.
• Overcome the limitations of prior polynomial approximation methods that rely on high-degree polynomials (e.g., 27-degree) leading to long multiplication chains and poor performance.
• Enable convergence of models using low-degree PAFs (8–14 degrees) through novel training techniques that stabilize optimization and preserve accuracy.
• Achieve a favorable trade-off between inference latency and model accuracy by systematically exploring low-degree PAFs with optimized training strategies.
• Eliminate reliance on hybrid schemes involving secure multi-party computation or garbled circuits by fully integrating PAFs into the FHE pipeline.

Proposed method

• Coefficient Tuning (CT): Adjusts PAF coefficients based on local input distributions before training to improve initialization and reduce accuracy drop.
• Progressive Approximation (PA): Replaces non-polynomial operators one layer at a time, enabling stable training via incremental fine-tuning under SGD.
• Alternate Training (AT): Decouples training between PAFs and other linear layers, using different hyperparameters to prevent interference and improve convergence.
• Dynamic Scale (DS) and Static Scale (SS): Dynamically scales PAF inputs to (−1, 1) during training and fixes the scale as the running maximum during FHE inference to prevent overflow.
• SMART-PAF framework orchestrates the sequential application of CT, PA, AT, DS/SS, along with dropout and SWA, to optimize the latency-accuracy trade-off.
• Employs a Pareto-frontier search strategy to identify optimal PAF degrees and configurations across models and datasets.

Figure 1: SMART-PAF replaces ReLU and MaxPooling by low-degree PAFs, and achieves better accuracy-latency Pareto-frontier than prior works Lee et al. ( 2021 ; 2022 ); Cheon et al. ( 2020 ) on ResNet-18 (ImageNet-1k).

Experimental results

Research questions

• RQ1 Can low-degree polynomial approximations (8–14 degrees) achieve high accuracy in FHE-based inference when replacing non-polynomial operators like ReLU and MaxPooling?
• RQ2 How can training divergence in high-degree PAFs be mitigated to enable convergence and preserve model accuracy?
• RQ3 What training techniques can stabilize optimization and improve convergence when combining linear layers with polynomial-approximated non-linearities?
• RQ4 Can a systematic framework identify the optimal trade-off between latency and accuracy in FHE inference across diverse models and datasets?
• RQ5 Does dynamic scaling during training and static scaling during FHE deployment effectively prevent numerical overflow while maintaining accuracy?

Key findings

• SMART-PAF achieves up to 13.64× higher accuracy than prior works on ResNet-18 with ImageNet-1k, significantly improving the accuracy-latency Pareto frontier.
• The framework enables a 14-degree PAF (f₁²∘g₁²) to achieve 69.4% post-replacement accuracy with 7.81× speedup compared to a 27-degree PAF from minimax approximation.
• For ResNet-18 on ImageNet-1k, SMART-PAF delivers 6.79× to 14.9× speedup and 1.42× to 13.64× accuracy improvement over prior methods.
• Coefficient Tuning improves initial accuracy by 1.04× to 2.38× across various PAF degrees without requiring fine-tuning.
• Progressive Approximation enables convergence of models with low-degree PAFs where standard training fails, particularly for degrees above 5.
• The combination of Dynamic Scale in training and Static Scale in deployment effectively prevents overflow and maintains numerical stability in FHE inference.

Figure 2: Overview of the FHE-base ML inference where original non-polynomial operators are replaced by Polynomial Approximated Activation (PAF).