Scalable privacy-preserving cancer type prediction with homomorphic encryption

Previous studies (deep_gene, 15, 16) on cancer detection using somatic mutations observed that the frequency of mutation of a gene correlates to higher prediction accuracy. The genes with higher number of mutations are more likely to develop a cancerous tumor. This makes sense because if there are more mutations in a gene, then the expression of that gene is likely disrupted through the mutations. There is also a correlation between the frequency of mutations in a patient cohort and the likelihood of observing the cancer in that cohort. This also makes sense because if a mutation is observed in a large number of patients with same cancer type diagnosis, then that mutation is likely either correlated or causal to that cancer type. Therefore, as the first step of feature selection, we choose genes with higher number of mutations. But each cancer type corresponds to a higher mutation in a different gene. We first rank the genes based on their SNV frequency in the patient cohort for each cancer type by also taking into consideration genes with more than one SNVs on them. We then combine the ranked genes from different cancer types and finally select the top 10,000 genes with highest SNV frequency for each cancer type as our features. This step also ensures removal of genes with low SNV frequency and reduce the dimensionality of our feature space. With all cancer types combined, we have a total of 18,606 genes.

Encoding techniques, both for the feature vectors and the target variables, ease in training towards better accuracy. Efficient encoding techniques have been proven to result in better performance in genomic classification tasks as well (encoding, 40). As mentioned in the dataset section, each SNV on a gene is represented by multiple characteristics. This information need to be meaningfully merged with the CNV information of each gene. We explore the following encoding schemes based on a biological intuition as explained below:

1. (1)

   Using presence of an SNV on a gene: The genes selected using frequency are merged for all cancer types. In this encoding we aim to study if a particular gene (mutation) is highly correlated to a cancer type. We assign a binary value $[0,1]$ to each of genes of a patient to denote the presence or absence of one or more SNVs. This allows us to encode SNV information in a categorical manner. This binary value per gene is used as a feature.

2. (2)

   Using the type and impact of an SNV: As denoted in subsection 3.1, the impact of mutation of a gene is calulated using VEP For each SNV in the dataset, we have two types of measure: 1- a qualitative measure indicating whether an SNV has a tolerated or deleterious effect; 2- a quantitative real-value measure ($s_{i,j}$) representing the strength of the impact and the qualitative confidence with which the mutation can be attributed to the diagnosis. $s_{i,j}$ is the strength of the $i^{th}$ SNV in patient $j$. Each of the encoding scheme are given equidistant real values between 0 to 1. We also experimented with different ranges but did not find any improvement in test accuracy. We think that both of these measures are useful in classifying the tumor type. We first encode the first measure by assigning values for [deleterious, deleterious (low confidence), tolerated (low confidence), tolerated] as [1.0, 0.75, 0.5, 0.25] for each $m_{i,j}$. The reason for this encoding is that a detrimental effect is given the highest value in cancer prediction and similarly a tolerated effect is given the lowest feature value. Either of the effects, when estimated with lower confidence are given lower effect. $m$ is the effect of the $i^{th}$ SNV in patient $j$. We then combine this encoding with the second measure as $m_{i,j}\times s_{i,j}$. The final effect values of SNV impact of a gene is the summation of the impacts of all the SNVs on that gene, if there is more than one SNVs. The resulting value per gene is used as a feature.

   In addition to the strength ($s$) of an SNV, the qualitative confidence of the effect of an SNV on a gene is also as a categorical variable with values as [high, moderate, modifier, low]. We encode them as [1, 0.4, 0.7, 0.1] since intuitively we want to assign a higher importance to a high mutation effect. A gene with no mutation is assigned 0 effect. Similar to the previous feature, for a gene with multiple SNVs at different locations, the values are added to finally represent the effective value of all the SNVs in a gene. The resulting value per gene is used as a feature (Figure 1).

3. (3)

   Using CNV of a gene CNV of a gene is represented as integers between -2 and 2, indicating if both, one, or no copy of the gene is deleted or duplicated. We scale these values to integers between 0 and 4 as statistical feature selection methods (for example $\chi^{2}$ test) often require positive values. Similarly, the resulting value per gene is used as a feature.

   Overall, the top 10,000 genes per cancer type are merged into a total of 18,606 genes. For each patient, 18,606 genes with their associated SNV encoding are concatenated with 25,128 genes with their CNV information. In total, we have 43,734 features which undergo the following statistical tests.

The previous steps of feature selection incorporate biological intuition. In this step, we explore statistical tests for evaluating feature importance. Statistical methods like $\chi^{2}$ test do not only inform the feature importance but also help reduce dimensionality of the feature space enabling faster computation (both during training and inference). It has previously been used in genetic information based disease prediction studies (chi_cancer, 41). We choose $\chi^{2}$ test as a feature selection metric since we achieve the best possible accuracy when compared to feature selection with mutual information or f-score statistics (as reported in Section 6). To decide if a feature is independent of the target label (i.e type of cancer), we perform the $\chi^{2}$ test where we calculate $\chi^{2}$ value of each feature with respect to the target variable. The $\chi^{2}$ value of a feature is given as $\sum\frac{(O_{i}-E_{i})^{2}}{E_{i}}$ where $E$ represents the expected value, $O$ represents the actual output and $i$ represents each instance of a $\chi^{2}$ test between a feature and a target. Note here that, the expected values of a variable is calculated using the distribution of feature values.

We run the $\chi^{2}$ test on all genes (CNV and SNV concatenated together) and sort the features decreasing order of $\chi^{2}$ values. The top $n$ features are selected and are used to train a classification model. We also use this step to analyze the selected genes and their relative importance in cancer type prediction in Section 6.1.5.

For the selection of the ML model that can correlate encoded somatic mutations to cancer type, we train a variety of ML models, while considering the difficulty of their implementation in HE domain. In plaintext, we experiment with Support Vector Machine (SVM) (with radial basis function, polynomial and linear kernels), logistic regression, and Deep Neural Networks (DNNs with two fully connected hidden layers with relu activation) possible classification models. We start with top 1,000 features selected by $\chi^{2}$ test and increase the number of features by 1,000 in each iteration. In our search for best-performing model, we train models using different number of features, different statistical tests for feature selection, different kernels (if applicable), with several regularization techniques, and with different optimization techniques cross-validated over 5 folds. We report the best performing models in Section 6.

Measures to reduce overfitting: We find that the logistic regression model performs best for several encoding schemes with the best test accuracy of $83.61\%$. In logistic regression, the probability that a sample $x$ belongs to class $k$ is given by $P(y=k|x)=\frac{e^{z_{k}}}{\sum^{K}_{l=1}e^{z_{l}}}$ where $z$ is the linear combination of coefficients of the form $\beta_{j}$. However, the number of features is much larger than the number of samples since each sample has effectively $\approx$ 43,000 features from CNV and SNV, whereas we have a total of 2,173 samples. Therefore, to avoid over-fitting in this high-dimensional setting, we introduce a Lasso ($l1$) penalty (lasso, 42) to the logistic loss function during training such that the features that are unlikely to contribute to the prediction are penalized and weighted zero. Therefore, if the logistic loss function is given by $L(\beta_{j})$ where $\beta_{j}$ represents the coefficients of the features, the loss function after Lasso penalty in the Lagrangian form becomes $L(\beta_{j})-\lambda\sum_{j}^{n}|\beta_{j}|$ which is minimized during training. Further, we use k-fold cross-validation on the training set ($k=5$) as another way to prevent the model from over-fitting to the training data. We iterate for 10,000 epochs to converge to a prediction model.

Metrics to detect problems of unbalanced data: High test accuracy on unbalanced datasets (with a higher percentage of samples from a particular label) can give a false sense of performance as a random guess (of the label with the highest number of samples) may also result in a high accuracy. For a holistic performance evaluation of our classifiers, we plot Receiver Operating Characteristics (ROC) Curve and report the individual area under curve for each class and the Micro-average Area Under Curve (MAUC) for the classifier. Although our dataset is unbalanced for 2 classes, still we report both test accuracy and ROCs for all the classes.

Binary models: Our cancer type prediction model uses the somatic genetic information to predict the cancer type from 11 classification labels. We also build models to predict each cancer type separately (like specific models in (gdl, 16)). These models are supplementary to our main prediction model and focus on one type of cancer. The features for this single cancer type model are chosen following the steps described above and the classifier is trained using a binary label: 0 for the all of the other cancer types and 1 for the cancer type of interest. We create separate models for each of the 11 cancer types and call these models as binary models since the prediction is converted into a binary classification task.

It should be noted here that the binary classification ability represented by the ROC curves of individual diseases (in our main prediction model) and the binary models for individual diseases are different because of the feature selection steps. In our binary models, the genes important to a specific disease are selected. However, for our main prediction model, the genes which are cumulatively important for all the 11 labels, are selected. Hence, we report both the analyses in section 6.

From analyzing the nature of genomic mutation data and the trends in accuracy (details in section 6), we observe that regardless of the ML model selected, the matrix multiplication would involve high-dimensional matrices (Dot product between weights and several thousands of features is common for all the ML models explored in our work). Therefore, following standard matrix multiplications would require a large number of multiplications corresponding to this high-dimensional dataset. The private cancer prediction methodology is characterized by a private inference protocol proposed in subsection 5.2.2 and then the fast matrix multiplication methodology, crux of the private ML algorithm, is proposed in subsection 5.2.3.

Fig. 2 shows the overview of our inference protocol. It consists of input encoding and encryption, weight and bias encoding, computation, decryption, and decoding. The client starts with a $|X|\times f$ matrix $X$ containing the input values represented with double-precision floating-point numbers, where $|X|$ is the number of inputs and $f$ is the number of features. The values of matrix $X$ are multiplied by a scaling factor $2^{s_{x}}$ in order to be converted into integers, a requirement of the BFV encryption scheme. This effectively converts our HE operations into fixed-point arithmetic. The scaled matrix of inputs $X_{s}$ is then encoded into a matrix of polynomial plaintexts $\bar{X}$, where each polynomial contains $n$ coefficients. We pack $n$ features of each row into a polynomial. This leads to an encoded matrix of dimensions $|X|\times\lceil f/n\rceil$.

It is worth noting that our model requires more precision than what can be represented in the plaintext modulus $t$. From experimental results, we determined that our inputs and weights require 14 bits of precision. Since our inputs are in the interval $0\leq x<2^{8}$, we set $s_{x}=6$ to represent the inputs in 14 bits. Meanwhile, weights are in the interval $0\leq w<1$, which leads to $s_{w}=14$. Due to the fixed-point arithmetic, the biases must be scaled by $2^{s_{x}+s_{w}}$. After the computation, the client will receive outputs that are scaled by a factor of $2^{s_{x}+s_{w}}$, like the biases, but that require $8+s_{x}+s_{w}+\lceil\log_{2}(f)\rceil$ bits of precision for representation. For $f=40960$, that translates to 44 bits. This is above of what a secure BFV ciphertext with enough noise budget for our computation can support. To cope with that without hindering the accuracy of our model, we used the Chinese Remainder Theorem (CRT) to break our plaintext into a pair of smaller plaintexts, each one under its own modulus. We define our plaintext moduli $T=\{t_{0},t_{1}\}$ as $t_{0}=1073872897$, which provides 30 bits of precision, and $t_{1}=114689$, offering 16 bits. This means that for every $n$ features encoded into a polynomial, we are actually encoding into a pair of polynomials, one with coefficient modulus $t_{0}$, and another with coefficient modulus $t_{1}$. For simplicity, we refer to this pair as plaintext polynomial.

The encoded matrix of inputs $\bar{X}$ is then encrypted with the client’s public key $pk$. The encrypted matrix $\hat{X}$ is sent to the server together with public values $\{n,T,s_{x}\}$. Afterwards, the server scales and encodes the transpose of the matrix of weights $W$, which has dimensions $f\times|Y|$, where $|Y|$ is the number of outputs. The transposition is a requirement of our computation. It packs several feature weights of an output in a plaintext polynomial, leading to an encoded matrix of weights $\bar{W}$ of dimensions $|Y|\times\lceil f/n\rceil$. Biases are encoded differently, each bias is encoded into a plaintext polynomial, filling all slots with its value. Finally, the server performs the matrix multiplication of encrypted inputs by encoded weights followed by addition of encoded biases $\hat{Y}=\hat{X}\times\bar{W}+\bar{B}$. The resulting matrix $\hat{Y}$ is returned to the client together with public value $s_{w}$. The client simply decrypts, decodes, and descale $\hat{Y}$ with its secret key $sk$ and obtains the result of the inference in plaintext.

Our privacy-preserving matrix multiplication algorithm, optimized for implementation in HE, is displayed in Algorithm 1. It receives three arguments: The encrypted matrix of inputs $\hat{X}$, encoded matrix of weights $\bar{W}$, and polynomial degree $n$. Each row of $\hat{X}$ represents an input, while each row of $\bar{W}$ represents one output. The computation of the dot product for each input is independent, making this algorithm highly parallelizable. We start the dot product by performing the column-wise multiplication of each row of $\hat{X}$ with the each row of $\bar{W}$ and append the result for each row-row pair into a vector (lines 5-11). Next, we add together all elements of the resulting vector (line 13) and execute $\log_{2}{n}$ ciphertext rotations and additions to finalize the dot product (lines 14-17). This results in a ciphertext where all its slots contain the result of the dot product of the row-row pair. In order to save memory and reduce communication time, we aim at packing several dot product results into a single ciphertext. For this, first we need to clear the ciphertext slots in all but one carefully chosen position. We do it by multiplying the resulting ciphertext $\hat{c}$ by a plaintext polynomial $\bar{p}$ with one at that specific position and zero in the remaining slots (lines 18-22). Finally, we can compress the dot product results $\hat{R_{0}}$ by adding them together (lines 25-42). If there are more dot product results than slots in a ciphertext, i.e., $|\hat{X}|\cdot|\bar{W}|>n$, then ciphertexts are appended to the output vector $\hat{Y}$. Lastly, we return the result $\hat{Y}$ (line 43). We provide the mathematical representation of the algorithm in Appendix

Tumor prediction is a classification problem which we address using multinomial Logistic Regression (LR). An LR model is trained by reducing the logistic loss function. During inference, the probability that an input $(x\in\mathbb{R}^{1\times d})$, with $d$ features, belongs to a class $(k)$ is given by $P(y=k|x)=\frac{e^{z}_{k}}{\sum^{K}_{l=1}e^{z}_{l}}$ where $z=Wx+b$, $W\in\mathbb{R}^{K\times d}$ is the weight matrix, and $b\in\mathbb{R}^{d}$ is the bias. The predicted class $(k_{p})$ is the class with the highest probability, i.e. $k_{p}=argmax({P(y=k_{i}|x)})$ where $k_{i}\in K$. This non-linear logistic function is computationally expensive in HE; thus, we perform the following approximation for building an ML model that can be used for encrypted inference. Since the logistic function is a monotonically increasing function, we can say that $P(y=k|x)$ for a class is higher if $z_{k}$ is higher, and since the predicted label depends on the relative probability values, the predicted label can also be calculated using $argmax({z_{k_{i}}})$. Therefore, during inference, a test input $(x_{test})$ needs to be multiplied with the weight matrix to get the final prediction, i.e. the predicted class $k_{p}=argmax(W\times x_{test}+b)$. Effectively, for efficient inference, the matrix multiplication between the test inputs and the weight matrix must be fast. Please note here that the size of $W$ is dependent on the number of features, i.e. the higher the dimension of an input, the larger is the $W$ matrix, and the more time-consuming is the matrix multiplication.

We report different performance metrics and the information on the features used for different models tested in Table 1. We observe that our model achieves a test accuracy of $66.85\%$ and a micro-average area under curve of $0.928$ with top 15,000 features. We also plot an Receiver Operating Characteristics (ROC) curve (Fig. 3) for each class and observe that skin cancer (class 9) detection has the highest area under the curve of $0.994$ while stomach cancer (class 10) prediction has the lowest area under the curve of $0.754$. Although on a slightly different dataset, other ML-based cancer prediction achieved a similar test accuracy of $65.5\%$ (deep_gene, 15) and of $70.08\%$ (gdl, 16). These methods also used the SNV frequency to prune the number of features.

We also experiment with just the copy number information for all the 25,128 genes and run a $\chi^{2}$ test to select the top genes. We achieve a slightly higher test accuracy of $71.27\%$ with 17,000 features. From Fig. 3, we observe that the micro average area also improves to $0.94$ with the lowest area under curve for detection of breast cancer (class 1) at $0.87$ (from $0.754$). Higher test accuracy and MAUC show that CNVs have more distinguishing power on the type of cancers than SNVs, when considered individually.

We select top $n$ features using both SNV and CNV information as described above and train several models to evaluate different machine learning algorithms for the tumor classification task. SVM with linear, rbf, and polynomial kernels achieve the best test accuracy of $68.13\%$, $64.82\%$, and $69.98\%$ with 13,000, 37,000, and 34,000 features respectively. Therefore, SVM does not show much improvement from our baseline that used only the presence of SNVs as features. Logistic regression model with Lasso penalty shows the best performance across all performance metrics. We achieve a test accuracy of $83.61\%$ with 34,000 features thereby also reducing the number of features. We also observe an improvement in the micro average area to $0.976$ when compared with models using just CNV or just presence of SNV. ROC for individual classes also have higher scores with the lowest area under curve of $0.94$ for cervical cancer (class 3). All the classes achieve an ROC area under curve of more than $0.9$. This experiment also shows that although CNVs are more informative, using both CNV and SNVs result in the highest prediction accuracy. We also test our logistic regression model using mutual information and f-score as feature selection methods but we achieve lower test accuracy of $81.95\%$ with 32,000 features, and $82.68\%$ with 40,000 features, respectively (Table1).

We built 11 specific models for each cancer type. We use both CNV and SNV information to select the top 34,000 features. In these experiments, we train 11 different models, where each one detects a particular tumor. We report the performance (test accuracy and micro-average score) in Table 3. We can see that all of the individual models have a test accuracy of more than $90\%$ and a micro-average area under curve of more than $0.9$. The best performing classifier is for kidney with a test accuracy of $97.97\%$ and a micro-average score of $0.996$. Even the worst performing binary model (for bronchus and lung) has a test accuracy of $91.34\%$ and micro-average score of $0.974$. Binary models have also been evaluated in other somatic genetic information-based cancer classification tasks (gdl, 16) where the authors also achieve high performance for individual tasks but suffer in performance in cancer prediction using all labels.

We discuss our findings on the top genes selected. 17,962 genes are selected for CNV data and 16,038 are selected for SNV data as the most informative features. However, more than $60\%$ of these genes are common, i.e. 11,133 genes are selected based on both CNV and SNV information. Considering the $\chi^{2}$ scores, we also observe that the top 1,030 genes come from using SNV information. We plot histograms of $\chi^{2}$ statistic scores in Fig. 4 and observe that the genes selected based on SNVs is more flat i.e. there are more genes with higher scores. We also verified that the highest score of a gene selected based on SNVs is $\approx 10\times$ higher than the highest value of a gene selected based on CNVs. Therefore, from feature selection perspective using $\chi^{2}$ test, SNV data on the selected genes are statistically more important than their CNV counterparts. But CNV information improves the test accuracy by adding potentially more relevant biological information.

When we investigated the top 10 most informative genes based on the SNV information (Table 4), we found “PTEN” are “APC” genes, which are known tumor suppressors; “MUC16” gene, which is a biomarker for ovarian cancer; “ZFHX3” gene, which is implicated in prostate cancer; “CCDC168” gene, which is known to be associated with Prostate Carcinoma and Uterine Body Mixed Cancer. The other 5 genes in this list are also implicated in important cellular activities that could potentially be related to cancer. These genes and their corresponding Gene Ontology (GO) term enrichment are depicted in Table 4.

The first gene in the top 10 most informative genes based on the CNV information (Table 4) is the first ever known tumor suppressor “RB1”. Similarly, “CDKN2A” and “CDKN2B” genes are also known tumor suppressors. “RCBTB2” gene is known to be repressed in prostate cancer. The “CDKN2B-AS1” gene has the silencing power of many other genes in the genome and strongly implicated in various cancer types. The other 5 genes in this list are also implicated in important cellular activities that could potentially be related to cancer.

The first insight is that biological intuition combined with high-dimensional data analysis methods can together achieve high accuracy (MAUC) while reducing the effective number of features. We further validate, using other studies, that our ML model indeed selects features that are biologically relevant. Using domain expertise, our ML model achieved a high performance of 0.98 MAUC with logistic regression ML model.

We further plot the variation of test accuracy (of a model trained using best hyper-parameters) with the number of features in Fig. 5 and observe that the test accuracy clearly rises when the number of features is lower than the number of samples and begins to saturate only after 30K features. From Fig. 5 we see that to achieve a test accuracy of more than $80\%$, we need at least 20K features, and to achieve the highest test accuracy, we need more than 30K features. Therefore, private genomic analysis is one such application where extremely high-dimensional must be processed. The second insight motivates the development of matrix multiplication algorithm, which when implemented using BFV, can result in fast yet private cancer prediction.

We report the encryption, decryption, and computation time required for private cancer prediction in Fig. 6. The time taken to calculate the final cancer label, which is effectively the result of matrix multiplication $(Wx_{test}+b)$, is denoted by computation. Computation, understandably, is the most costly operation in private cancer prediction. We observe that even if the number of features increase from 16K to more than 40K ($2.5\times$), the computation time only increases from 33.44 seconds to 35.52 seconds ($1.06\times$), which corresponds to $\approx 7\%$ increase in test accuracy. Therefore, the matrix multiplication is not the bottleneck for private cancer prediction. The time needed for encryption of the test samples increases with the number of features, with 3.87 seconds for 16K features to 10.40 seconds for 40K features ($2.68\times$) which indicates a linear increase in the encryption time as a function of number of features. Decryption is the least expensive operation (less than 1 second) as compared to encryption and computation; the values for decryption time are labelled in Fig. 6. The maximum total time for private inference of the entire test dataset is required when processing 40,960 features, and is 46.77 seconds.

As mentioned in section 1 private computations using HE are generally designed for high throughput, since popular FHE schemes support batching. For our application, we also prioritize latency, i.e. evaluation of a single sample. We report our findings in Fig. 7. From the figure we observe that the total amount of time to privately compute cancer label for a sample is 1.08 seconds and there is a linear increase in time with the number of samples.

In order to accurately quantify the performance benefits of our proposed methodology, we implement a privacy-preserving version of standard matrix multiplication using BFV. For the experiment we generate synthetic data in the same format (CNVs and SNVs) for 8192 individuals and measure time for encryption, computation, decryption operations. We plot the timing results in a log-log graph in Fig. 8. We observe that the total time required for private inference implemented using standard matrix multiplication for similar number of individuals as the test set is approximately 10 minutes, approximately $10\times$ more than our methodology. Also, the total time required for private inference on 1 individual is 598.25 seconds (similar time required for thousands of individuals), which is $550\times$ more than the time required by our algorithm. Therefore, as compared to standard matrix multiplication, commonly used for implementation of ML models, our algorithm has lower latency, and higher throughput.

Healthcare models are difficult to port trivially across datasets (as discussed in section 1). Cancer detection ML model is no exception. However, our matrix multiplication algorithm is not dependent on input data or weight values (like quantization-based DNN design techniques (qnn, 45)) and thus, can be reused for datasets requiring HE-based high-dimensional inference. The transferability of our private inference algorithm across applications is an added advantage.