Efficient Privacy-Preserving KAN Inference Using Homomorphic Encryption

KAN, as a promising alternative to MLP, has been integrated with popular models by various researchers. In (Zhang and Zhang 2024), Zhang et al. combined KAN with Graph Neural Networks, utilizing KAN for feature extraction and experimentally demonstrating the effectiveness of their proposed approach. Xu et al. (Xu et al. 2024) introduced KAN to enhance the message-passing process in Graph Collaborative Filtering, proposing a more efficient recommendation model. In (Genet and Inzirillo 2024), Genet et al. incorporated Temporal KAN into the Temporal Fusion Transformer to further improve the model’s ability to handle time series data. Li et al. (Li et al. 2024) combined KAN with U-Net for the segmentation and generation of medical images. However, these works primarily focus on the performance of KAN within neural networks and do not address the issues of security and privacy leakage in practical applications.

In 2016, Ran et al. introduced CryptoNets (Gilad-Bachrach et al. 2016), achieving privacy-preserving inference based on HE, though with very low efficiency. Since then, researchers have continued to improve inference performance on convolutional networks (Chou et al. 2018; Brutzkus, Gilad-Bachrach, and Elisha 2019; Benaissa et al. 2021), though the network structures they implemented are relatively simple. Both Lee et al.(Lee et al. 2022) and Ran et al.(Ran et al. 2023) employ more efficient ciphertext packing methods to achieve privacy-preserving inference for complex convolutional networks. Zimerman et al. (Zimerman et al. 2024) achieved privacy-preserving inference for Transformers. Additionally, researchers have combined Graph Convolutional Networks (GCN) with HE to enhance privacy protection for cloud-based graph data inference (Ran et al. 2022; Peng et al. 2024; Ran et al. 2024). There are also many privacy-preserving inference schemes based on MPC (Juvekar, Vaikuntanathan, and Chandrakasan 2018; Srinivasan, Akshayaram, and Ada 2019; Hao et al. 2022; Pang et al. 2024), but the communication overhead of MPC-based schemes is enormous, making them unsuitable for resource-constrained users. KANs differs significantly from current deep models. The aforementioned methods rely heavily on specific model structures, making them unsuitable for direct application to KAN privacy-preserving inference. Therefore, it is necessary to design a new method to achieve privacy-preserving inference for KAN.

Unlike traditional neural networks that use fixed activation functions, KANs employs learnable activation functions at the network’s edges, enabling each weight parameter to be replaced by a univariate function. However, identifying suitable basis functions can be challenging, as many cannot be represented effectively. To address this, KANs utilizes spline functions for parametric approximation, providing significant flexibility. This method allows for modeling complex functions with fewer parameters, thereby enhancing the model’s interpretability.

The flexibility of spline functions enables them to adaptively model complex relationships in data by adjusting their shape, minimizing approximation errors, and improving the network’s ability to learn subtle patterns from high-dimensional datasets.

In most cases, $\operatorname{spline}(x)$ is parametrized as a linear combination of B-splines as follows (Liu et al. 2024):

where $c_{m}$ are trainable parameters, and $B_{m,k}(x)$ are the $k$-order B-spline basis functions at the knots. The $k$-order B-spline basis functions are defined recursively as follows:

where $t_{1},\cdots,t_{n+k}$ form a non-decreasing sequence of knots, which determine the domain and influence of the basis functions.

The overall network in KAN can be expressed by the following equations:

Following current work in privacy-preserving inference (Lee et al. 2022; Ran et al. 2023), we use RNS-CKKS (Cheon et al. 2019) as the underlying cryptographic scheme. Below, we briefly introduce the main operations in HE.

RNS-CKKS encrypts a vector at a time using Single Instruction, Multiple Data (SIMD). To clearly describe the operations of HE, we use bold lowercase letters to represent vectors, and $\llbracket\cdot\rrbracket$ to denote encrypted ciphertexts. For example, we use $\mathbf{x}$ to represent the vector $(x_{1},x_{2},\cdots,x_{n})$, and $\llbracket\mathbf{x}\rrbracket$ to denote its ciphertext. We use $\oplus$, $\ominus$, $\otimes$, and $\operatorname{rot}($) represent homomorphic ciphertext addition, subtraction, multiplication, and rotation.

where $t$ is the rotation step. When $t>0$, the ciphertext rotates left; when $t<0$, it rotates right. A more detailed theoretical analysis and security assessment can be found in (Cheon et al. 2017, 2019).

The threat model in this paper is similar to that of previous HE-based privacy-preserving inference schemes (Lee et al. 2022; Ran et al. 2023). We assume a cloud-based machine learning service scenario where the cloud server hosts a pre-trained KAN model with plaintext weights. Clients send their sensitive data to the cloud server to obtain inference results from the network model. We assume the cloud server is honest-but-curious, meaning it follows the protocol but may attempt to extract users’ private information through additional computations. To protect the privacy of the client’s sensitive data, the client can encrypt the data using HE before sending it to the cloud server for privacy-preserving inference. The cloud server performs computations on the ciphertext without decrypting it, and the client decrypts the inference results using their private key. This process ensures that no information about the input data is revealed to the cloud.

Previous approximations of the ReLU function utilized the sign function, which allows the input to be scaled to the $[-1,1]$ range, resulting in a good approximation. However, SiLU cannot be approximated using sign or comparison functions. The biggest challenge in approximating SiLU is determining the appropriate approximation range.

The input range of the activation function varies across different datasets (Lee et al. 2023). We estimate the mean $\mu$ and variance $\sigma^{2}$ of the activation function’s input range by feeding the dataset into a pre-trained network. Then, we use Chebyshev’s inequality to estimate the approximate distribution of the activation function’s input values.

According to Chebyshev’s inequality,

when we set $\varepsilon=5\sigma$, we assume that approximately $96\%$ of the data falls within the $[(\mu-5\sigma,\mu+5\sigma)]$ interval. Therefore, we use $[\max(\mu-5\sigma,x_{\text{min}}),\min(\mu+5\sigma,x_{\text{max}})]$ as the approximation range to ensure that most data lies within this interval.

Next, we uniformly sample points within the approximation range and use weighted least squares to fit the activation function. We focus on fitting the regions with dense data distribution and de-emphasize the sparse edge regions. This approach allows us to obtain a polynomial that better approximates the original activation function under the same polynomial degree constraints. Let the input be $x_{i}$, the output of the activation function be $y_{i}$, and the output of the polynomial be $f(x_{i})$. The optimization objective of weighted least squares is:

where $n$ is the number of sampled points, and $w_{i}$ is the weight. The closer the data is to the mean, the higher the weight we assign. For example, the sum of squared residuals of data points within the $[\mu-3\sigma,\mu+3\sigma]$ interval is given a weight $w_{i}$ of 10, while those outside this interval are given a weight of 1. We then solve for the polynomial coefficients to obtain the approximate polynomial for the activation function.

We observed that the calculation of B-spline basis functions involves many parallelizable components. Therefore, we use Repeat Packing to fully leverage the SIMD capabilities of HE to accelerate B-spline calculations.

Let the input be a vector $\mathbf{x}$ of length $n_{i}$. After repeat packing, it becomes $\llbracket\mathbf{x}|\mathbf{x}|\mathbf{x}\dots|\mathbf{x}\rrbracket$, denoted as $\llbracket\mathbf{x}\rrbracket_{\text{rep}}$, where $\mathbf{x}$ is repeated $g+2k$ times, with $g$ being the number of B-spline grids and $k$ being the B-spline degrees, satisfying $n_{i}\cdot(g+2k)\leq\frac{N}{2}$.

Repeat packing data in ciphertexts is not easy, as ciphertext rotations are computationally expensive. We designed a fast packing method that reduces the number of rotations needed for repeat packing from $g+2k$ to $\lceil\log_{2}(g+2k)\rceil$. We first rotate $\llbracket\mathbf{x}\rrbracket$ to the right by $n_{i}$ steps, then add it to the original $\llbracket\mathbf{x}\rrbracket$ to obtain a temporary ciphertext $\llbracket\mathbf{x}|\mathbf{x}\rrbracket$. The above steps are repeated on the temporary ciphertext, with rotation steps of $n_{i},2n_{i},4n_{i},\cdots$. The specific algorithm is shown in Algorithm 1.

To simplify the description, we define the following operations:

• •

  $\mathtt{ColTile}$$(\mathbf{A},l,r)$: Concatenates columns $l$ to $r-1$ of matrix $\mathbf{A}$ into a vector.

• •

  $\mathtt{PolyComp}$$(a,b)$: Returns $1$ if $a>b$, $0$ if $a<b$, and $\frac{1}{2}$ if $a=b$. We implement this using the method from (Lee et al. 2021).

Recall Equation 2 and Equation 3. We first need to determine whether the input $\mathbf{x}$ falls within the interval $[t_{m},t_{m+1}]$, and then recursively compute the basis functions.

We combine all the intervals into a weight matrix $\mathbf{G}$. Then, using $\mathtt{ColTile}$, we obtain a vector $\mathbf{g_{1}}$ containing the left endpoints of all intervals $[t_{m},t_{m+1}]$ and a vector $\mathbf{g_{2}}$ containing the right endpoints. Next, we use $\mathtt{PolyComp}$ to compare the packed input $\llbracket\mathbf{x}\rrbracket_{\text{rep}}$ with $\mathbf{g_{1}}$ and $\mathbf{g_{2}}$, and then recursively calculate the basis functions according to Equation 3.

It should be noted that $\mathtt{PolyComp}$ operates within the range $[-1,1]$. If the input exceeds this range, the approximation may become inaccurate. Therefore, we scale the input to the $[-1,1]$ range. We estimate the approximate range of the input using the training data, and represent the input within a larger range $[-R,R]$ to avoid overflow. When using $\mathtt{PolyComp}$ to compute intervals, there may be cases where $x$ is exactly equal to the interval endpoint, such as $x=t_{m}$. In this case, $\mathtt{PolyComp}$ returns $\frac{1}{2}$, which does not match Equation 2. We can achieve the correct result by comparing $2x$ with $2t_{m}+1$. However, the probability of equality is very small, so its impact on the overall approximation error is negligible. Therefore, in practice, we can ignore the equality case, similar to how existing papers approximate ReLU (Lee et al. 2022). The specific algorithm is shown in Algorithm 2.

In the previous section, we described how to approximate B-spline basis functions. However, according to Equation 1, we need to multiply each basis function by a coefficient $c_{m}$ and sum them to obtain the final B-spline function approximation. For a given input $x$, this is an inner product between the coefficient vector $\mathbf{c}$ and the basis function vector $\mathbf{b}$. Since we are using SIMD, computing the inner product requires rotations, which are expensive in the HE domain. Therefore, we adopt a lazy combination approach, postponing the computation of $\operatorname{spline(x)}$ until later to reduce the number of rotations. Specifically, we combine $\mathbf{c}$ with the weight coefficients $\mathbf{W_{s}}$ of the subsequent linear layer to obtain a new set of weights. This reduces the number of multiplications and rotations, improving computational efficiency.

The weight matrix $\mathbf{W_{s}}$ has dimensions $n_{o}\times n_{i}$. When the input size is $n_{i}$, the coefficients of the basis functions form a matrix $\mathbf{C}$ of size $n_{i}\times(g+k)$. After combining $\mathbf{W_{s}}$ and $\mathbf{C}$, we obtain a large matrix $\mathbf{W^{\prime}}$ of size $n_{o}\times n_{i}\cdot(g+k)$. Therefore, following the method of (Blealtan 2024), we directly use $\mathbf{W^{\prime}}$ during inference instead of the linear layer weights $\mathbf{W_{s}}$ and $\mathbf{C}$.

However, due to the repeat encoding introduced during B-spline basis function computation, the result is in $\mathtt{ColTile}$ format, with an output vector of length $n_{i}\cdot(g+k)$. We need to repack the result of the basis function approximation into row-major order to take advantage of SIMD for fast matrix operations. This is essentially a ciphertext reordering problem, requiring multiplication by a permutation matrix $\mathbf{P}$ of size $n_{i}\cdot(g+k)\times n_{i}\cdot(g+k)$. The method for generating the permutation matrix is provided in Algorithm 3.

To reduce multiplicative depth and the number of rotations, lazy combination is also used as mentioned above. We multiply the permutation matrix $\mathbf{P}$ with the weight coefficients $\mathbf{W^{\prime}}$. The final fusion linear layer weight is then given by:

The computation of the fusion linear layer, as well as the linear layer following SiLU, involves matrix multiplication. We optimize this using the baby-step giant-step (BSGS) strategy (Halevi and Shoup 2021), reducing the number of ciphertext rotations required for matrix multiplication and lowering the time cost of the linear layer.

Compared to the naive method without lazy combination optimization, our proposed B-spline basis approximation method offers a significant efficiency improvement. Table 3 demonstrates that our method improves computational efficiency several times compared to the naive method. For example, when $n_{i}=256$, $g=10$, and $k=5$, our method takes only $42.241$ seconds, which is $5.03\times$ faster than the naive method. This improvement is because the naive method requires a permutation operation after B-spline basis approximation to rearrange the data. The permutation in HE domain introduces additional costs, including $n_{i}\cdot(g+k)$ HE ciphertext multiplications and $\sqrt{n_{i}\cdot(g+k)}$ HE rotations, leading to significant latency overhead. As the input dimension $n_{i}$, B-spline grid $g$, and B-spline degree $k$ increase, the speedup of our method improves.

Additionally, the results in Table 3 show that, as the parameters $n_{i}$ and $g$ increase, the HE inference latency remains relatively stable with our method. This is because our method is designed to take full advantage of SIMD of CKKS for parallel processing. Therefore, as long as the single ciphertext packing condition $n_{i}\cdot(g+2k)\leq\frac{N}{2}$ is satisfied, increases in $n_{i}$ and $g$ do not significantly affect the HE inference latency. However, as parameter $k$ increases, we needs to iteratively compute higher-order B-spline basis functions, increasing the HE inference latency due to the added multiplicative depth. In practical applications, the performance of the B-spline basis approximation can be optimized by carefully selecting appropriate values for $n_{i}$, $g$, and $k$.

To evaluate the practical performance of our proposed privacy-preserving KAN inference scheme, we conducted experiments on various datasets. We compared with plaintext KAN as a baseline, as well as with MLP (Rumelhart, Hinton, and Williams 1986), which is a fundamental component in current HE-based privacy-preserving inference methods (Benaissa et al. 2021; Lee et al. 2022; Ran et al. 2023). Apart from image classification tasks, we also conducted experiments on symbolic formulas. We selected the following symbolic functions from (Liu et al. 2024) for our experiments. The Toy Formula dataset represents a high-dimensional symbolic function $f(x_{1},\cdots,x_{100})=\exp(\frac{1}{100}\sum_{i=1}^{100}\sin^{2}(\frac{\pi x_{i}}{2}))$, while the lpmv0 and lpmv1 datasets consist of special functions from (Liu et al. 2024).

The experimental results show that the accuracy of our privacy-preserving KAN inference is very close to that of the plaintext domain and exceeds that of MLPs, as shown in Table 4. On image benchmark datasets, the accuracy decrease of our scheme is no more than $0.25\%$, and it outperforms MLPs. On the symbolic formula dataset, the RMSE deviation from plaintext is no more than $0.00011$, with a lower RMSE than MLPs, confirming the effectiveness of our HE-based KAN inference scheme.

We evaluated the inference latency for a single input across different datasets in a single-threaded environment, comparing our method with the naive implementation, as shown in the figure. For the small dataset Toy-formula, our method achieves an inference latency of 324 seconds, while the naive method has a latency of 468 seconds, resulting in approximately $1.44\times$ speedup. On the larger CIFAR-10 dataset, our inference latency is 380 seconds compared to 2789 seconds for the naive method, yielding over $7\times$ speedup. Additionally, since servers often handle data from multiple users, the amortized processing time for multiple datasets is also important. Due to page limitation, we provide details on the multi-threaded environment’s amortized inference latency and the time distribution of various computational components in the Appendix.