Non-polynomial operators pose a challenge in FHE due to the lack of native support. To overcome this, previous research has explored (1) a hybrid scheme, which offloads non-polynomial operators to other secure schemes with a secure data transfer to communicate data among schemes, and (2) approximation, which replaces non-polynomial operators with polynomial approximation functions (PAF).

The hybrid scheme is challenging in practice because a plethora of prior arts illustrate the prohibitive communication overheads for transferring data securely among different schemes Gilad-Bachrach et al. (2016); Lou et al. (2021); Ran et al. (2022).

Despite its promise, PAF approximation is non-trivial because it requires striking a balance between accuracy and latency. A high-accuracy approximation requires high-degree PAFs with a prohibitively long chain of multiplications with bootstrapping. An example includes a SotA 27-degree PAF Lee et al. (2021); Kim et al. (2022). On the other hand, a low-latency approximation suffers from severe accuracy degradation, allowing only a subset of non-polynomial operators to be replaced with PAFs. In such cases, there are still some non-polynomial operators being offloaded to other schemes to reduce the accuracy drop, resulting in dominating communication latency overheads in practice Mishra et al. (2020a); Lou et al. (2021); Ran et al. (2022). These approaches are suboptimal and highlight the need for exploring alternative approaches.

To address this challenge, previous research has explored coefficients fine-tuning to reduce accuracy drop when replacing non-polynomial operators with low-latency low-degree PAFs. However, existing techniques fail to converge for PAFs with degrees higher than 5, and PAFs with degrees lower than 5 still suffer from severe accuracy degradation for a simple 7-layer CNN model under CiFar-10 dataset Lou et al. (2021), let alone deeper ResNet-18 model under more complex task like ImageNet-1k. To explore the optimal degree with minimal accuracy degradation for various models under different tasks, the key challenge is to improve training techniques to enable the convergence of PAF with arbitrary degrees.

To tackle the aforementioned dilemma, this paper proposes four key techniques and a framework to enable the convergence of the PAF-approximated model when using PAFs with arbitrary degrees. This is the first-ever framework, to the best of our knowledge, to (1) replace all non-polynomial operators with $8\sim 14$-degree PAFs, and (2) to adopt proposed training techniques to minimize accuracy degradation. Out contributions are listed as follows. (The comparison with prior arts is shown in Tab. 1)

Pre-Training Novel Techniques:

• $\bullet$

  Coefficient Tuning (CT): Using a uniform PAF to replace all non-polynomial layers neglects variations in input distributions, causing a marked drop in accuracy of PAF-approximated model (Tab. 3). CT refines PAF coefficients based on local input distributions, enhancing the accuracy of PAF-approximated model without fine-tuning by $1.04\times\sim 2.38\times$ across PAFs with various degrees.

• $\bullet$

  Progressive Approximation (PA): Directly replacing all non-polynomial operators and training the full network, as adopted by previous works, lacks a theoretical convergence guarantee under SGD (Fig. 9). Instead, PA sequentially replaces non-polynomial operators in inference order, with each substitution followed by training layers preceding the replacement point. The theoretical convergence of PA under SGD is analyzed in §3.1, resulting in an accuracy boost of $0.4\%\sim 1.9\%$.

In-Training Novel Techniques:

• $\bullet$

  Alternate Training (AT): When training PAF and linear operator (like convolution) in a PAF-approximated model together, issues arise: slow convergence with a small learning rate or divergence with a large one. This stems from the different hyperparameter needs of PAF coefficients and linear layers training. Therefore, AT trains PAF coefficients and other layers separately, using different hyperparameters, and alternates between PAF coefficients training and other layers training. This leads to accuracy climbing by $0.6\%\sim 2.3\%$ (Tab. 3).

• $\bullet$

  Dynamic Scaling (DS) and Static Scaling (SS): PAF training is sensitive to the input value range. Prior works take a fixed scale value to reduce the input range in training because FHE does not support value-dependent operators. However, input values in training might overflow under a fixed scale, hampering fine-tuning accuracy and potentially leading to divergence. We address this by introducing Dynamic Scaling, normalizing inputs to the $[-1,1]$ range during training. Subsequently, we apply Static Scaling to the post-training model, anchoring the scale to the peak value in FHE deployment.

SMART-PAF Framework:

• $\bullet$

  The order of applying proposed techniques affects final accuracy, and thus a systematic scheduling framework is proposed to automatically perform non-polynomial replacement and PAF-approximated model training. Our evaluation results show that for ResNet-18 (ImageNet-1k) and VGG-19 (CiFar-10), SMART-PAF consistently enables low-degree PAFs to demonstrate higher accuracy than high-degree SotA PAFs (Tab. 3 and Tab. 4). For example, SMART-PAF identifies the Pareto-frontier of various PAFs shown in Fig. 1 and spots the sweet-point PAF with 14 degrees, which can achieve the same 69.4% validation accuracy with $7.81\times$ latency speedup in ResNet-18 inference on ImageNet-1k, compared to the minimax approximation with 27-degree PAF Lee et al. (2021).

This work pushes the state-of-the-art in PAF-approximated model training for higher accuracy with lower latency (Fig. 1). First, we formalize PAF-approximated model training into an optimization problem (§3.1). Second, §4 elaborates on the SMART-PAF techniques and framework with both intuitive and theoretical analysis. Third, we individually assess the SMART-PAF techniques for VGG-19 (CiFar-10) and ResNet-18 (ImageNet-1k) in §5.

FHE is an asymmetric encryption scheme that enables ciphertext-based computation with Cheon-Kin-Kim-Song (CKKS) Cheon et al. (2016) as the most commonly used FHE scheme for machine learning inference due to its superior efficiency in approximate computation compared to other schemes such as BGV, BFV, and TFHE Riazi et al. (2020). Under the CKKS scheme, only polynomial operators are allowed as shown in Fig. 2 such that all non-polynomial operators including both ReLU and MaxPooling must be replaced by PAF (Approximation) or offloaded to other schemes (Hybrid Scheme). Prior works have also shown that approximation offers superior performance and lower overhead compared to the hybrid scheme Lou et al. (2021).

High-degree polynomials are theoretically capable of approximating arbitrary functions, but approximating ReLU or MaxPooling directly can result in severe approximation errors Lee et al. (2022). Prior research Lee et al. (2022) has shown that it is more effective to approximate the $sign(x)$ function ¹¹1$sign(x)$ outputs $1/-1$ if $x$ is positive/negative, and $0$ for zero and then construct the ReLU and Max operator using it, such as with $\frac{(x+sign(x)\cdot x)}{2}$ and $\frac{(x+y)+(x-y)\cdot sign(x-y)}{2}$.

Lower latency and higher accuracy (low accuracy degradation) are two fundamental goals of replacing non-polynomial operators with PAFs. Latency is determined by the degree of the polynomial, as FHE-based multiplication dominates latency. Accuracy degradation arises from the difference between the PAF and the original non-polynomial operators, as PAF cannot precisely approximate targeted non-polynomial functions. Such differences can vary depending on the degree of the PAF, the cascaded format of the PAF, and the coefficient values of the PAF.

We adopt cascaded polynomial based PAFs throughout the paper because they achieve lower approximation error than a single polynomial with the same degree Lee et al. (2022; 2021). Tab. 2 shows the selected PAFs with the minimal multiplication depth under different degrees constraints. We use the notation $f^{n}$ to indicate a serial nested function call of the same polynomial $f$ for $n$ times, for example, $f^{2}=f(f(x))$.

The effectiveness of PAFs in approximating non-polynomial operators depends on their input range. For example, setting the approximation input range as $[-1,1]$ for a PAF indicates that it gives relatively more accurate approximations of non-polynomial operators’ output when the input falls in $[-1,1]$ than outside this range.

However, determining an appropriate input range can be challenging because a narrow input range can lead to severe approximation loss for input value outside the input range and potentially lead to training divergence. In contrast, a broader range may result in a large average approximation error across the entire input range, because of the limited representation capability of PAFs to capture negligible differences among input values under a broad range.

Prior works Lee et al. (2021); Lou et al. (2021) adopted a fixed scale in training because FHE does not support value-dependent operation. Specifically,  Lee et al. (2021) determines the fixed scale by adding some margin to the maximum of input values. However, input values in training might overflow the range specified by a fixed scale, hampering fine-tuning accuracy and potentially leading to divergence. Therefore, we propose dynamic-scale training and only pick a fixed scale in model deployment.

The essential problem for approximating non-polynomial operators with PAF is the regression problem shown in Eq. 1. Specifically, the goal is to find a vector of coefficients $\mathbf{a}_{i}=\{a_{i}^{0},\cdots,a_{i}^{N}\}$ for a N-degree PAF, such that the cumulative error of replacing non-polynomial function $R(\mathbf{x}_{i})$ by a $N$-degree PAF is minimal given input data as $\mathbf{x}_{i}\ (i=0,\cdots,D-1)$, where $D$ stands for the total number of non-polynomial layers ²²2we use non-polynomial layers and non-polynomial operators interchangeably. Boyd et al. (2004).

$\min\limits_{\mathbf{a}_{0},\cdots,\mathbf{a}_{D-1}}f(\mathbf{a},\mathbf{x})=$ $\sum_{i=0}^{D-1}\left[\frac{1}{N}\sum_{i=0}^{N-1}(R(\mathbf{x}_{i},\mathbf{a}_{0},\cdots,\mathbf{a}_{i})-\mathbf{a}_{i}\cdot\mathbf{x}_{i})^{2}\right]$

(1)

where $\mathbf{x}_{i}=\{x_{i}^{0},x_{i}^{1},\cdots,x_{i}^{N}\}$. The subscript / superscript of $x$ and $a$ indicates layer index / degree index. On the other hand, the subscript / superscript for $\mathbf{a}$ and $\mathbf{x}$ indicates layer index / the training epoch. After replacing $0,\cdots,(i-2)$-th ReLU by PAFs, the input of the $i$-th ReLU $\mathbf{x}_{i}$ becomes dependent on $\mathbf{a}_{0},\cdots,\mathbf{a}_{i-1}$, i.e. $i$-th ReLU becomes $R(\mathbf{x}_{i},\mathbf{a}_{0},\cdots,\mathbf{a}_{i})$ instead of $R(\mathbf{x}_{i})$.

The Eq. 1 is non-convex as it optimizes all $D$ vectors of coefficients $\mathbf{a_{i}},\ i\in[0,D)$ under the changing $\mathbf{x}$.

$$\min\limits_{\mathbf{a}_{i}}f^{\prime}(\mathbf{a}_{i})=\\ \frac{1}{N}\sum_{i=0}^{N-1}(R(\mathbf{x}_{i})-\mathbf{a}_{i}\cdot\mathbf{x}_{i})^{2}$$

(2)

The Eq. 2 only optimizes a single layer, such that $\mathbf{x}_{i}$ statistically is fixed, which makes it convex. In such cases, the stochastic gradient descent could guarantee finding the optimal solution $\mathbf{a}_{i}^{*}$ given initial value as $\mathbf{a}_{i}^{0}$. The error could converge in the speed order as shown in Eq. 3 after training $i$-th non-polynomial operator coefficients $\mathbf{a}_{i}$ for $T$ epochs Boyd et al. (2004) $i\in[0,D)$. And a good initialization value of $\mathbf{a}_{i}^{0}$ could reduce the total error and thus improve the convergence speed.

$$f^{\prime}(\mathbf{a}_{i}^{T})-f^{\prime}(\mathbf{a}_{i}^{*})=O(\frac{||\mathbf{a}_{i}^{0}-\mathbf{a}_{i}^{*}||^{2}}{\sqrt{T}})$$

(3)

where $\mathbf{a}_{i}^{T}$ indicates the value of $\mathbf{a}_{i}$ after training for $T$ epochs.

The PAF-approximated model replaces non-polynomial operators with PAFs. Such a replacement changes the original model structure, leading to changed data distribution of $x_{i}$ for all layers after the replacement points. Therefore, original coefficients trained using the original model structure may perform worse under new data distribution, leading to accuracy degradation. Therefore, the PAF replacement forces retraining to fine-tune coefficients to learn the changes in input data distributions.

Training methods adopted by prior arts lead to divergence when training PAF with a degree higher than 5, forbidding exploration using PAFs with higher degrees to enhance overall accuracy. In this section, we analyze critical reasons behind the divergence of existing training algorithms and propose corresponding four techniques to guarantee training convergence while replacing non-polynomial operators with PAFs of arbitrary degrees. Then we propose a framework to schedule techniques automatically.

Eq. 3 indicates that a good initialization can reduce the overall training time, facilitating faster convergence, which has been ignored by prior arts Lee et al. (2021); Cheon et al. (2020); Lee et al. (2022); Lou et al. (2021). These studies have typically initialized PAFs with the same coefficients using traditional regression algorithms. Such initialization is suboptimal because it overlooks the distribution differences of data at different non-polynomial operators.

Therefore, we propose coefficient tuning (CT) as a technique to use profiled data distribution ³³3We assume that distribution probability range is consistent across input samples. to obtain a closer-to-original initialization point and then use PAF with different initialization points to replace different non-polynomial operators, as shown in Fig. 3. The CT involves the following steps: (1) obtain the coefficients of PAFs using traditional regression methods on the full range of input data; (2) profile the input data distribution for each non-polynomial operator; (3) tuned PAF coefficients to minimize overall approximation errors on profiled distributions to obtain the post-CT PAFs; (4) replace original non-polynomial operators with the post-CT PAFs.

Intuitively, CT tunes PAFs to achieve higher accuracy and less approximation error over a reduced smaller input range, which corresponds to the highest-probability range in the data distribution. Besides, CT reduces training time because the initial value of coefficients $\mathbf{a^{0}}$ are closer to the optimal value $\mathbf{a^{*}}$ as indicated by Eq. 3.

Replacing all non-polynomial operators in the given ML model with PAFs simultaneously, as done by prior arts, is a suboptimal approach. This is because approximation errors of early replacements propagate to all layers behind the replacement point, making the regression (Eq. 2) non-convex by varying both PAF coefficients $\mathbf{a}_{i}$ and input data $\mathbf{x}_{i}$ ($i\in[0,D)$) for all $D$ non-polynomial layers, causing training divergence. Instead, we propose the Progressive Approximation (PA) approach, which replaces the non-polynomial layer with PAF, one layer at a time, followed by fine-tuning of PAF coefficients until accuracy convergence as shown in Fig. 4. In such cases, all $\mathbf{x}_{i}$ ($i\in[0,D)$) statistically get fixed because we only vary PAF coefficients of a single layer, which is one of all ($\mathbf{a}_{0},\mathbf{a}_{1},\cdots,\mathbf{a}_{D-1}$), ensuring a simple convex regression problem shown in Eq. 2, which is easily optimizable by SGD.

For example in Fig. 4, PA starts with replacing the first ReLU with a PAF, and then fine-tunes PAF coefficients until accuracy converges. Then PA repeats the same flow for the following non-polynomial operators.

Intuitively, the replacement of all non-polynomial operators by PAFs simultaneously introduces a huge deviation to the original model which is hard for the training algorithm to recover, while PA applies the overall approximation error progressively instead of applying all errors at once, restricting the approximation error to the optimizable range of the training algorithm to enable convergence and effectively mitigate the accuracy degradation.

Replacing the non-polynomial operators with PAFs requires model fine-tuning, as pretrained model parameters are no longer optimal for the new PAF-approximated model. However, multiple previous works fine-tune the PAF-approximated models as a whole to optimize both PAF coefficients and the parameters of other layers (like Convolution, BatchNorm, etc.). Such a training scheme, however, deteriorates accuracy, indicating a failure of convergence. We demonstrate it later in Fig. 9. This is because modifications to PAF coefficients can have a significant impact on final inference results, unlike changes to convolution weights, which may not affect overall results at all. Such an observation is consistent with the model structure that all data need to go through PAF while some weights can be pruned without any impact on overall accuracy.

To tackle training divergence, PAF coefficients fine-tuning should be decoupled from fine-tuning parameters of other layers, and two fine-tuning processes should use different training hyperparameters. We thus propose Alternate Training (AT) to fine-tune PAF coefficients and parameters of other layers separately in an alternate manner. Specifically, AT fine-tunes PAF coefficients first while keeping other parameters fixed, as shown in AT step 1 in Fig. 5. After reaching a specific epoch threshold, PAF coefficients get frozen and AT fine-tunes parameters of other layers in AT step 2. Note that the training in different AT steps may use different hyperparameters because of different parameter sensitivity in different layers. AT repeats until the accuracy finally converges, often resulting in an accuracy climb.

Intuitively, AT considers the sensitivity difference between PAFs parameters and parameters of linear layers, and decouples the two training processes to avoid training interference.

In Section 2.3, we discuss the challenge of setting an appropriate input range for PAFs. Previous studies have employed a large input range, such as $[-50,50]$, to prevent infinite errors for values outside the range Lee et al. (2021). However, we contend that this approach is suboptimal because it results in high approximation error across the entire input range, which can impede training convergence and lead to significant accuracy degradation (e.g., accuracy drops from 69% to 10% for ResNet-18 on ImageNet-1k).

To tackle this issue, we introduce Dynamic Scaling (DS) and add an auxiliary layer before each PAF to normalize the value range of its input data to $[-1,1]$ during fine-tuning. DS determines the scale on a batch-by-batch basis. For each batch, the scale is set to the highest absolute input value, and then all data are proportionally scaled to fit within the $[-1,1]$ range. This ensures that the inputs within each batch spread the $[-1,1]$ range, maximizing their relative differences for better distinguishment while staying within bounds.

However, the post-training model is intended for FHE deployment, where Dynamic Scaling (DS) is inapplicable because FHE cannot select the batch’s maximal value due to the absence of value-dependent operators. Therefore, we propose Static Scaling (SS) to fix the scale of each auxiliary layer in PAF-approximated model as the input running maximum under the training dataset Such a scale determination relies on the similarity between training data and validation data. In our evaluation, either a higher or smaller scale results in lower accuracy for both VGG-19 (CiFar-10 dataset) and ResNet-18 (ImageNet-1k dataset).

The sequence in which CT, PA, AT, and DS/SS are applied can greatly affect the final validation accuracy. Furthermore, employing existing techniques, such as dropout for overfitting mitigation and Stochastic Weights Averaging (SWA) for faster convergence, is crucial.

To navigate these complexities, we introduce a scheduler in SMART-PAF. This scheduler systematically applies training configurations and proposed techniques, organizing the training into steps, with each step replacing a single non-polynomial layer with PAF. Coefficient Tuning is applied offline before all steps. The sequence within a single step is outlined below:

• $\bullet$

  Training Group: A training group first trains the PAF-approximated model for $E$ epochs, followed by applying SWA using weights of all $E$ epochs. The model with the highest validation accuracy proceeds to the following procedures.

• $\bullet$

  Accuracy Improvement Detection: If the validation accuracy gets improved during the process, the framework launches a new training group until no accuracy improvement is observed.

• $\bullet$

  Overfitting Avoidance: The framework applies Dropout followed by a new training group if spots overfitting ⁴⁴4we adopt the empirical overfitting condition, which is “training accuracy $>$ validation accuracy + $10\%$”.

• $\bullet$

  Alternative Training: When previous techniques do not improve accuracy anymore, AT is applied to swap training targets between PAF and other layers, followed by a new training group.

• $\bullet$

  Step Termination Condition: Current step terminates when observing no accuracy improvement in above processes.

SMART-PAF framework dedicates one step for each non-polynomial layer, orchestrating these steps progressively based on the inference order (Progressive Approximation). Dynamic Scaling is adopted in the entire fine-tuning process, while Static Scaling is applied for the post-finetuning model.

Model We evaluated SMART-PAF with models being used by prior works for fair comparison Lee et al. (2021); Lou et al. (2021), including VGG-19 (18 ReLU and 5 MaxPooling), and ResNet-18 (17 ReLU and 1 MaxPooling).

Datasets We adopt CiFar-10 and ImageNet-1k He et al. (2015); Krizhevsky et al. (2012) to test the generality of proposed techniques for tasks with different complexity.

PAF Form We adopt 6 PAFs (Tab. 2) with minimal multiplication depth under different degrees as the approximation of $sign(x)$. ReLU and Max operators are replaced by $\frac{(x+sign(x)\cdot x)}{2}$ and $\frac{(x+y)+(x-y)\cdot sign(x-y)}{2}$, separately.

Latency Evaluation We implement PAF in Microsoft SEAL library Chen et al. (2017) using CKKS Cheon et al. (2016) (Degree: 32768, modulus bitwidth: 881) and evaluate the wall-clock PAF latency on AMD Threadripper 2990WX.

Accuracy Evaluation The accuracy is obtained through evaluating PAF-approximated models under given datasets.

CT adjusts PAF coefficients based on profiled value distributions and improve the $1.05\sim 3.32\times$ post-replacement validation accuracy w/o fine tuning, as shown in Fig. 7.

When only approximating ReLU, CT shows more benefits for polynomials with lower degrees, as they have less capability to fit the entire input range and are more prone to significant accuracy reduction. CT mitigates this loss by focusing on fitting the high-probability region in the distribution, resulting in less initial post-replacement approximated error. On the other hand, polynomials with higher degrees have less overall approximated error across the entire input range and, therefore, show less improvement from CT.

Further, approximating both ReLU and MaxPooling leads to $10.9\%\sim 21\%$ accuracy drop than approximating ReLU only, as shown by comparing the blue bar in the top and bottom figures in Fig. 7. Contrary to the ReLU approximations in earlier works Park et al. (2022), this suggests that solely assessing ReLU does not provide a full picture of PAF’s non-polynomial approximation prowess. Comprehensive evaluations should encompass both ReLU and MaxPooling, wherein CT consistently aids in accuracy restoration, improving up to $34.1\%$ accuracy for $f_{1}^{2}\circ g_{1}^{2}$!

During fine-tuning, the baseline (prior works) uses a direct replacement + direct training approach: it replaces all non-polynomial operators with PAFs and trains other layers, excluding the PAFs. In contrast, PA adopts a step-wise strategy. In each step, it replaces one non-polynomial operator with a PAF and trains preceding layers progressively (termed progressive replacement and progressive training). This method yields an accuracy improvement ranging from $0.4\sim 1.7$%, as illustrated in Fig. 8. Significantly, the progressive replacement is pivotal for this improvement. This is evident from the marked accuracy degradation of the green bar (direct replacement + progressive training) relative to the orange bar. This is because progressive replacement simplifies the optimization, enabling its theoretical convergence under SGD, as discussed in §3.1.

While PA doesn’t always ensure superior accuracy—for instance, the baseline training (blue) outperforms PA in the $f_{1}\circ g_{2}$ setup—it can enhance post-fine-tuning accuracy beyond the limits of higher-degree PAFs shown in Fig. 8. Consequently, the SMART-PAF framework divides training into smaller steps, proceeding with the best results between PA and baseline training, as depicted in Fig. 6.