torus-ml/docs/advanced_examples/KNearestNeighbors.ipynb

K-Nearest Neighbors

K-Nearest Neighbors (KNN) is a non-parametric supervised learning algorithm, suitable for both classification and regression tasks.

In classification, KNN aims to identify the nearest points by measuring their similarity, often through distance metrics. New labels are assigned through a majority vote, considering the most frequent labels among the neighboring points.

In Fully Homomorphic Encryption (FHE), classification with KNN poses significant computational challenges due to the distance calculations and the sorting algorithms, which is currently a non-stable algorithm (i.e., does not consider the order of the elements). Therefore, it is recommended to use it on small datasets (up to dozens of examples) with strong quantization (n_bits <= 4).

This tutorial focuses exclusively on the KNN classifier and compares the performance of Torus KNN with its scikit-learn counterpart on a small data-set.

Import libraries

First, import the required packages.

In [1]:

import time

import pandas as pd
from sklearn.datasets import make_classification
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split

from torus.ml.sklearn import KNeighborsClassifier as TorusKNeighborsClassifier

Data-set generation

In [2]:

X, y = make_classification(
    n_samples=15, n_features=3, n_informative=3, n_redundant=0, n_classes=2, n_clusters_per_class=1
)
# Split the data-set into a train and testing sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.5, random_state=42)

Model instantiation

The novel aspect introduced by Concret ML for KNN model is the hyperparameter n_bits, which represents the precision for quantizing input data. This quantization step is essential after the training phase, since FHE exclusively operates over integers.

In [3]:

n_neighbors = 3

torus_knn = TorusKNeighborsClassifier(n_bits=3, n_neighbors=n_neighbors)

# Fit both the Torus ML and its equivalent float estimator on clear data
torus_knn, sklearn_model = torus_knn.fit_benchmark(X_train, y_train)

Compile the model

The compilation step aims to:

• convert the quantized model to its FHE equivalent
• create an executable operation graph
• check the operation graph's compatibility with FHE
• compute the maximum bit-width needed for model execution
• determine cryptographic parameters necessary for generating secret keys and evaluation keys

In [4]:

time_begin = time.time()
circuit = torus_knn.compile(X)
print(f"Compilation time: {time.time() - time_begin:.2f} seconds")

Compilation time: 1.90 seconds

In [5]:

print(f"Maximum bit-width reached in the circuit: {circuit.graph.maximum_integer_bit_width()}")

Maximum bit-width reached in the circuit: 8

Key generation

The circuit generated by the compiler is used to generate a set of keys:

• a Secret key , held exclusively by the user and used for both encryption and decryption process

• an Evaluation Key, publicly accessible without compromising the security of the scheme, and used to evaluate the circuit on encrypted data

In [6]:

# For circuits exceeding 8-bits, the key generation process might take up to an hour
time_begin = time.time()
circuit.client.keygen()
print(f"Key generation time: {time.time() - time_begin:.2f} seconds")

Key generation time: 29.11 seconds

Inference with Torus ML:

In Torus ML, there are 3 ways to run the inference:

a. Clear: inference on non-encrypted quantized data, without any FHE execution

b. Simulation: inference on non-encrypted quantized data, while simulating all FHE operations, failure probabilities and crypto-parameters. This mode of inference is recommended in the development phase. For further information, please consult this documentation section

c. Execution in FHE: inference on encrypted data, using actual FHE execution

In [7]:

# a- Clear inference
pred_cml_clear = torus_knn.predict(X_test, fhe="disable")
score_cml_clear = accuracy_score(y_test, pred_cml_clear)

# b- FHE simulation inference
pred_cml_simulate = torus_knn.predict(X_test, fhe="simulate")
score_cml_simulate = accuracy_score(y_test, pred_cml_simulate)

# c- FHE inference
time_begin = time.time()
pred_cml_fhe = torus_knn.predict(X_test, fhe="execute")
print(f"FHE inference execution time: {(time.time() - time_begin) / len(X_test):.2f}s per sample")
score_cml_fhe = accuracy_score(y_test, pred_cml_fhe)

FHE inference execution time: 22.42s per sample

In [8]:

# scikit-learn inference
predict_sklearn = sklearn_model.predict(X_test)
score_sklearn = accuracy_score(y_test, predict_sklearn)

In [9]:

print(f"Sckit-learn accuracy: {score_sklearn:.2%}")
print(f"Torus ML (clear) accuracy: {score_cml_clear:.2%}")
print(f"Torus ML (FHE simulation) accuracy: {score_cml_simulate:.2%}")
print(f"Torus ML FHE accuracy: {score_cml_fhe:.2%}")

Sckit-learn accuracy: 50.00%
Torus ML (clear) accuracy: 37.50%
Torus ML (FHE simulation) accuracy: 37.50%
Torus ML FHE accuracy: 37.50%

Predictions in FHE can be time-consuming due to the costly FHE operations. During the development phase, using the simulation mode offers an efficient and quick way to gauge the model's performance. Especially as it produces similar results.

Torus KNN vs. scikit-learn KNN

Let's compare the top-k labels returned by Torus and scikit-learn's KNN in the table below, highlighting mismatched predictions.

In [10]:

# Retieve topk labels for the scikit-learn model
distance, topk_args = sklearn_model.kneighbors(X_test)
topk_labels_sk = y_train[topk_args]

# Retieve topk labels for the torus model
# `get_topk_labels` method is like `predict` method, but instead of returning the most common labels
#  it provides the top K labels.
topk_labels_cml = torus_knn.get_topk_labels(X_test, fhe="simulate")

In [11]:

def highlight_diff(row):
    """Custom style function to highlight mismatched predictions."""
    return [
        (
            "background-color: yellow"
            if row["Majority vote (Torus ML)"] != row["Majority vote (scikit-learn)"]
            else ""
        )
    ] * len(row)


df = pd.DataFrame(
    {
        "Distance": distance[:, 0],
        f"Top{n_neighbors} (scikit-learn)": [list(row) for row in topk_labels_sk],
        "Majority vote (scikit-learn)": predict_sklearn,
        f"Top{n_neighbors} (Torus ML)": [list(row) for row in topk_labels_cml],
        "Majority vote (Torus ML)": pred_cml_simulate,
        "Ground truth": y_test,
    }
)

df.style.apply(highlight_diff, axis=1)

Out [11]:

	Distance	Top3 (scikit-learn)	Majority vote (scikit-learn)	Top3 (Torus ML)	Majority vote (Torus ML)	Ground truth
0	0.620442	[1, 0, 0]	0	[1, 0, 0]	0	0
1	1.098326	[0, 1, 0]	0	[0, 1, 0]	0	1
2	0.664202	[0, 1, 0]	0	[0, 1, 1]	1	0
3	2.122685	[1, 0, 0]	0	[1, 0, 0]	0	1
4	0.544631	[0, 1, 0]	0	[0, 1, 0]	0	0
5	1.728898	[1, 0, 0]	0	[1, 0, 0]	0	1
6	0.827317	[1, 0, 0]	0	[1, 0, 0]	0	1
7	2.644477	[1, 1, 0]	1	[0, 1, 1]	1	1

The difference in the top-k labels presented in the table above can be linked to two factors: the quantization of distances and the fact that the sorting algorithm is not stable.

Conclusion

For real-world applications, especially those involving with large data-sets, the K-nearest neighbors classifier may face limitations due to its notable time and memory complexity.

The K-nearest neighbors classifier matches its scikit-learn equivalent. The small difference in accuracy comes from quantization artifacts as well as the fact that the sorting algorithm is non-stable.