Automated machine learning for secure key rate in discrete-modulated continuous-variable quantum key distribution

In recent decades, machine learning (ML) has gained impressive breakthroughs that deeply impact both industry and academia, including autonomous driving Grigorescu et al. (2020); Levinson et al. (2011), natural language processing Deng et al. (2013); Young et al. (2018), protein structure prediction Jumper et al. (2021) and even proving mathematical conjectures Davies et al. (2021). ML aims to recognize patterns in data, especially multidimensional data, and generalize them to new instances, which contributes to automating tasks and reveals hidden patterns beyond humans intuition. This modern information-processing technology also benefits solving intractable quantum tasks, since quantum tasks are usually counterintuitive and involve high dimensions. Several significant advances have been made by applying ML to quantum physics, from classifying quantum states Gao et al. (2018); Yang et al. (2019); Ahmed et al. (2021), quantum control Bukov et al. (2018); Lumino et al. (2018); Niu et al. (2019) to quantum metrology Hentschel and Sanders (2011).

Quantum key distribution (QKD) enables unconditional security between two legitimate users (Alice and Bob) against any eavesdropper called Eve Bennett and Brassard (1984); Ekert (1991), which is guaranteed by quantum mechanics laws Shor and Preskill (2000). According to different detection methods, QKD is currently divided into two categories: discrete-variable (DV) QKD Lo et al. (2014); Gisin et al. (2002) and continuous-variable (CV) QKD Grosshans and Grangier (2002); Lance et al. (2005); Huang et al. (2015); Yin et al. (2019). Between these two categories, CV QKD has unique edges on a higher secret key rate and excellent compatibility with standard communication components Fossier et al. (2009); Huang et al. (2016); Jouguet et al. (2012), which enables CV QKD to be competitive at a metropolitan distance Pirandola et al. (2020). To enhance the practicality of CV QKD, several works introduce machine learning-based methodologies to the CV QKD area, such as developing a novel CV QKD scheme Jin et al. (2021); Liao et al. (2020), parameter prediction Liu et al. (2018) and detecting quantum attacks Mao et al. (2020).

CV QKD protocols with discrete modulation have attracted increasing attention for decades. Its appealing advantages include easier experimental implementation and higher error-correction efficiencies which promote CV QKD over longer distances Xu et al. (2020); Leverrier and Grangier (2009, 2011); Zhao et al. (2009). These properties bring potential advantages in large-scale deployment in quantum-secured networks Simon (2017). However, the security analysis of discrete-modulated CV QKD protocols is more complicated owing to the lack of symmetry Coles et al. (2016). Recently, some novel numerical approaches Lin et al. (2019); Winick et al. (2018) have been proposed to analyze the security of discrete-modulation protocols against collective attacks, where key rate calculation involves minimizing a convex function over all eavesdropping attacks that are consistent with the experimental data. These numerical approaches achieve much higher key rates over significantly longer distances compared with previous security analyses. Based on these numerical approaches, a neural network model was presented to quickly predict the secure key rate of discrete-modulated CV QKD with high reliability (secure probability as high as $99.2\%$). This neural network model learns the mapping between input parameters and key rates from datasets generated by numerical methods, which supports the computation of secure key rates in real time Zhou et al. (2021). However, the mapping complexity between input parameters and key rates depends on the solving complexity of discrete-modulated protocols’ key rates through numerical approaches Hu et al. (2021). Selecting architectures and hyperparameters plays a critical role in the performance of a neural network. Therefore, to learn different mappings from different protocols, the architectures of neural networks and the corresponding hyperparameters should be adjusted carefully by humans, which comes at a great price Yu and Zhu (2020).

Here, we propose a more flexible and automatic neural network model combined with Bayesian optimization Shahriari et al. (2015), which maintains extremely high reliability and efficiency and reduces complicated manual adjustment. Our method is universal for a variety of unstructured QKD protocols that lack analytical tools and rely on numerical methods. We apply our model to two variants of discrete-modulated CV QKD protocols and acquire high secure key rates with considerable tightness in both cases. We then compare the time consumption of our model with the numerical method proposed in Ref. Lin et al. (2019), which shows a great speedup of approximately $10^{7}$.

This paper is organized as follows. In section  II, we introduce the numerical method for CV QKD with discrete modulation proposed in Ref. Lin et al. (2019), and we rely on it to collect a dataset to train and test the model. In section  III, we introduce more details about the Bayesian optimization used in this paper. In section  IV, we demonstrate all the main results of this paper. Section  V provides a discussion and concludes this paper.

In this work, we apply the model in two discrete-modulated CV QKD protocols with different detection techniques to demonstrate the generalizability of our model. One is the quadrature phase-shift-keying (QPSK) heterodyne detection protocol Lin et al. (2019), and the other is an improved QPSK homodyne detection protocol Liu et al. (2021). To collect a dataset for training neural networks, we generate secure key rates of both protocols by applying the same numerical method Lin et al. (2019); Hu et al. (2021). In the following, we briefly introduce how computing key rates can be transformed into a relevant convex objective function for numerical optimization. A more detailed description can be found in Ref.Lin et al. (2019).

Here, we consider a CV QKD protocol with quaternary modulation that involves two parties: a sender Alice and a receiver Bob. During each time in an iteration of N rounds, Alice randomly prepares one of the four coherent states $\left|\alpha_{k}\right\rangle=\left||\alpha|e^{i(2k\pi/4+\pi/4)}\right\rangle$, where $k\in\{0,1,2,3\}$, and sends it to Bob via an untrusted quantum channel. Then, Bob uses either homodyne or heterodyne detection to estimate k. The secret key rate under collective attacks in the asymptotic limit is given by the following expression according to the Devetak-Winter formula Devetak and Winter (2005)

where $H\left(\mathbf{Z}\mid E\right)$ is conditional von Neumann entropy, which describes the uncertainty of the string $\mathbf{Z}$ in Eve’s view. Eve’s maximal knowledge of Bob’s string $\mathbf{Z}$ requires the minimum uncertainty of $\mathbf{Z}$ under a certain density matrix $\rho$. Therefore, we need to find the optimum $\rho^{*}$ in feasible domain $\mathbf{S}$ to minimize $H\left(\mathbf{Z}\mid E\right)$, $p_{\text{pass }}$ is the sifting probability, and $\delta_{\mathrm{EC}}$ is the actual amount of information leakage per signal in the error-correction step. To turn this problem into a convex optimization problem, the above expression can be reformulated as

in which $D(\rho\|\sigma)=\operatorname{Tr}\left(\rho\log_{2}\rho\right)-\operatorname{Tr}\left(\rho\log_{2}\sigma\right)$. As shown in Ref.Winick et al. (2018), $\mathcal{G}$ is a completely positive and trace nonincreasing map that describes the postprocessing of different quadratures. $\mathcal{Z}$ is a pinching quantum channel that reads out the key information.

Since the term $p_{\text{pass }}\delta_{\mathrm{EC}}$ in formula $\left(2\right)$ is easy to compute, we can only consider the following relevant optimization problem:

where $k\in\left\{0,1,2,3\right\}$, $\langle\hat{q}\rangle_{k}$, $\langle\hat{p}\rangle_{k}$, $\langle\hat{n}\rangle_{k}$ and $\langle\hat{d}\rangle_{k}$ denote the expectation values of corresponding operators when Bob measures states labeled by $k$. These expectation values can be obtained through homodyne or heterodyne measurements. These first four constraints come from experimental outcomes. The next two constraints are natural requirements since $\rho_{AB}$ is a density matrix. The last constraint on the partial trace of system B comes from the fact that the quantum channel cannot influence system A of Alice. We can handle the above density matrix and operators in finite dimensions $N_{c}$ after imposing the photon-number cutoff assumption on this optimization problem Ghorai et al. (2019); Lin et al. (2019). Then, this problem can be solved numerically. Eventually, we solve this minimization problem by the numerical method proposed in Ref.Winick et al. (2018). The specific implementation of this numerical method in our work can be found in Ref.Liu et al. (2021). This method involves two steps:

1. Find a solution that is close to optimal, which gives an upper bound on the key rate.

2. Convert this upper bound to a lower bound on the key rate by considering its dual problem.

In this section, we present a brief introduction to Bayesian optimization. Bayesian optimization is a powerful strategy for global optimization of objective functions that are expensive to evaluate Shahriari et al. (2015); Bergstra et al. (2011). This method is gaining great popularity in hyperparameter optimization. In particular, hyperparameter optimization in machine learning can be represented as follows:

where $f(x):\mathcal{X}\rightarrow\mathbb{R}$ is an objective function to minimize, $x^{\star}$ is a hyperparameter vector yielding the lowest value of $f$, and the dimension of domain $\mathcal{X}$ depends on the total type of concerned hyperparameters. In practice, the evaluation of the objective function is extremely costly, which leads to selecting proper hyperparameters by hand becoming intractable. Beyond the manual tuning method, grid search and random search Bergstra and Bengio (2012) are two common methods that perform slightly better. However, these methods still waste a large amount of time evaluating poor hyperparameters across the entire search space, which is relatively inefficient. In contrast, Bayesian optimization estimates the true objective function with a probability model. Then, it utilizes Bayes’ theorem to update this model based on previous results and chooses the next promising hyperparameters. In practice, this method can find better hyperparameters in less time. Figure 1 illustrates the Bayesian optimization procedure.

Sequential model-based optimization (SMBO) algorithms are formalizations of Bayesian optimization. Bergstra et al. (2011) These algorithms have two key ingredients:

1. A probabilistic surrogate model $S$. SMBO approximates the objective function $f$ with a probabilistic model called a surrogate, which is cheaper to evaluate. This surrogate contains a prior distribution capturing beliefs about the behavior of the objective function and is then updated sequentially after each new trial.

2. An acquisition function $\alpha:\mathcal{X}\rightarrow\mathbb{R}$. The acquisition function is the criterion by which the next vector of hyperparameters is chosen from the surrogate function.

For an SMBO algorithm at iteration n, the next location $x_{n+1}$ is selected by optimizing $\alpha_{n}$ and to evaluate the true $f$ to obtain a result $y_{n+1}=f(x_{n+1})$. The new tuple $(x_{n+1},y_{n+1})$ is appended to the historical set $\mathcal{H}$. Then, the surrogate model $S$ is updated incorporating the new results, which means that the prior is updated to produce a more informative posterior distribution over the space of objective functions. The pseudocode of this framework is summarized in Algorithm 1.

The most common choice of acquisition function is expected improvement (EI):

Here $y^{*}$ is a threshold value of the objective function $f$, and $p_{S}\left(y\mid x\right)$ represents the surrogate probability model. If this expectation is positive, then the vector of hyprparameters $x$ is expected to produce a better result than $y^{*}$. There are several different strategies for constructing the surrogate model: a Gaussian process approach Williams and Rasmussen (2006), random forests Breiman (2001) and a tree-structured Parzen estimator(TPE) Bergstra et al. (2011). In this work, the TPE approach is adopted, which supports continuous, categorical and conditional parameters, as well as priors for each hyperparameter over which values are expected to perform best Hutter et al. (2015). In contrast, the Gaussian process approach and random forests only support one or two types of the above parameters, which are not capable of our following task covering continuous, categorical and conditional parameters. Instead of directly modeling $p\left(y\mid x\right)$, this method models $p\left(x\mid y\right)$ using two such densities over the configuration space $\mathcal{X}$:

This algorithm chooses $y^{*}$ to be some quantile $\gamma$ of the observed y values, which means $p\left(y<y^{*}\right)=\gamma$. So the $EI_{y^{*}}(x)=\frac{\gamma y^{*}\ell(x)-\ell(x)\int_{-\infty}^{y^{*}}p(y)dy}{\gamma\ell(x)+(1-\gamma)g(x)}\propto\left(\gamma+\frac{g(x)}{\ell(x)}(1-\gamma)\right)^{-1}$. The tree-structured form of $\ell$ and $g$ makes it easy to draw many candidates according to $g(x)/\ell(x)$. On each iteration, the algorithm returns the candidate $x$ with the greatest EI. We implement this algorithm for the hyperparameter optimization of the neural networks predicting CV QKD key rates, by using a Python library called Hyperopt Bergstra et al. (2013).

Artificial neural networks can approximate arbitrary bounded continuous mapping on a given domain, according to universal approximation theorem Hornik et al. (1989). Therefore, we expect that the neural network can learn the mapping between input variables defined in the constraints of Eq.(3) and output key rates, which avoids solving the time-consuming optimization problem and computes key rates with low latency. We demonstrated this possibility of using a neural network to predict the key rates of discrete-modulated CV QKD in previous work Zhou et al. (2021). In that work, we built a four-layer fully connected forward neural network holding a loss function designed specifically to predict the key rates of discrete-modulated CV QKD with homodyne detection. The objective loss function is the key ingredient to keep the output key rates reliable and tight. We retain it in this work but utilize the TPE algorithm to search other parts of the neural network to improve the network’s overall performance. The specific formula of the loss function is as follows:

For training inputs $\left\{\vec{x}_{i}\right\}$ and corresponding labels $\left\{{y}_{i}\right\}$, here $n$ is the size of $\left\{\vec{x}_{i}\right\}$, $e_{i}^{*}=y_{i}^{*p}-y_{i}^{*}$ is the residual error between preprocessed label $y_{i}^{*}$ and the corresponding output of the neural network $y_{i}^{*p}$, where $y_{i}^{*}=-\log_{10}\left(y_{i}\right)$. There are two significant hyperparameters $\gamma$ and $\varepsilon$ contained in this loss function, the choices of which are crucial to a model’s performance, as we presented in Ref.Zhou et al. (2021). The meaning of hyperparameter $\gamma$ is to force the predicted key rate as information-theoretically secure as possible, and $\varepsilon$ is to force the predicted key rate as close to numerical results as possible. Here, apart from the input layer and output layer, we do not fix the structure of the neural network but utilize the TPE algorithm to search it efficiently in a set configuration space. An illustration of our model is shown in Fig. 2.

After this training under TPE searching is complete, we obtain the resulting structures of neural networks in both cases, which are shown in Table 3-4. Then, we use the selected and trained network to predict key rates on the test set for both protocols. The predicted key rates that show security achieve as high as $99.15\%$ for the QPSK heterodyne detection protocol and $99.59\%$ for the QPSK homodyne detection protocol, which suggests that our method combining a neural network with Bayesian optimization is highly reliable. For those key rates predicted securely, namely, predicted results are lower than the true values, we plot their relative deviation distributions for both protocols in Fig. 3. Figure 3 suggests that our method has good tightness.

Before training the neural network under the TPE method, we generate datasets for two different protocols by the aforementioned numerical approach. To obtain datasets with diversity, for the QPSK heterodyne detection protocol, we generate $36$ sets of data from excess noise $\xi=0.0045-0.0405$. Each dataset contains $80$ random samplings for $\xi$ from an interval of length $0.001$, for example $\left[0.0045,0.0055\right]$. Under each random sampling, we generate data every $5$km with the transmission distance $L$ up to $200$km. At each distance, we generate data from amplitude $\alpha=0.62-0.72$ in a step of $0.01$. The total datasets contain $809,600$ input instances $\left\{\vec{x}_{i}\right\}$ and corresponding labels $\left\{{y}_{i}\right\}$. For the QPSK homodyne detection protocol, excess noise is sampled randomly from $\xi=0.014-0.042$, where the length of sampling interval is $0.002$, for example $\left[0.014,0.016\right]$, and amplitude $\alpha$ is sampled from $\left[0.60,1.05\right]$. The size of total datasets is $368,116$. For both protocols, each $\vec{x}_{i}\in\left\{\vec{x}_{i}\right\}$ represents a vector of $29$ variables, there are $16$ variables that are the right parts of the first four constraints of Eq. 3, $12$ variables are nondiagonal elements of the right side matrix of the last constraint of Eq. 3, and the remaining variable is excess noise $\xi$. Label $y_{i}\in\left\{\vec{y}_{i}\right\}$ represents the corresponding key rate. Therefore, we fix the neurons of the network’s input layer in $29$ and output layer in $1$, and the search space of other hyperparameters can be found in Table 1-2.

Before feeding data into neural networks, we split data into a training set and a test set and implement data preprocessing as in Ref.Zhou et al. (2021). For the QPSK heterodyne detection protocol, the training set contains $769,120$ data instances, and the test set contains $40,480$ data instances. For the QPSK homodyne detection protocol, the training set contains $327,636$ data instances, and the test set contains $17,244$ data instances. For both cases, there is $10\%$ of the training data split as the validation set. We generate the dataset on the blade cluster system of the High Performance Computing Center of Nanjing University. We consume over $250,000$ core hours, and the node we use contains 4 Intel Xeon Gold 6248 CPUs, which involves immense computational power. Under the TPE algorithm with max iteration $10$, the Adam algorithm Kingma and Ba (2014) is used to train neural networks for $200$ epochs, and the initial learning rate is set to $0.001$. It takes roughly $53$ hours for the QPSK heterodyne detection protocol and $23$ hours for the QPSK homodyne detection protocol on an Nvidia A100 GPU.

Here, we also compare the predicted results with numerical results in key rates versus transmission distance for two protocols. The comparison is shown in Fig. 4. For this plot, we implement the same numerical approach to compute the best key rates of two protocols for different excess noises by optimizing the amplitude $\alpha$ of signal states in the range $\left[0.62,0.72\right]$ and $\left[0.62,1.03\right]$ with a step of $0.01$. The choice of the excess noise range is consistent with the sampling interval of previous training data. The photon-number cutoff $N_{c}$ is $12$, and the maximal iteration number of the first step in the numerical approach $N_{i}=300$. We record the corresponding $29$ variables producing the best key rates as neural networks’ inputs to predict key rates. As shown in Figs. 4(c) and (d), the predicted results are all secure and remain tight with relative deviations between $10\%$ and $20\%$ when the transmission distance is below $150$ km for both protocols.

To show the efficiency of our method, we compare the running time between the neural network method and the numerical method on a high-performance personal computer with a 3.3 GHz AMD Ryzen 9 4900H 16 GB of RAM, as shown in Fig. 5. The results suggest that the neural network method is generally 6-8 orders of magnitude of the numerical methods. For example, when $\xi=0.025$, the numerical method consumes approximately $850$ seconds to calculate the key rate at $50$ km for the QPSK heterodyne detection protocol. When $\xi=0.035$, the numerical method consumes approximately $1260$ seconds at $25$ km to calculate the key rate for the QPSK homodyne detection protocol. However, we can use a trained neural network to obtain results in approximately $0.0001$ seconds, which is almost real time.

To summarize, we develop a neural network model combined with Bayesian optimization to directly extract key rates with high reliability, considerable tightness and great efficiency. Beyond designing the neural network architecture by human and troublesome manual tuning of hyperparameters, we utilize a special Bayesian optimization method called the TPE algorithm to automatically search the structure and hyperparameters that are the best fit for a given dataset. We exemplify our method on two promising discrete-modulated CV QKD protocols varied by different detection techniques across a large range of excess noise and transmission distances. For both protocols, the neural networks selected by the TPE algorithm predict the information-theoretically secure key rates with great high probability(up to $99.15\%$ for the QPSK heterodyne detection protocol and $99.59\%$ for the QPSK homodyne detection protocol), and the results present considerable tightness.

We show that our method can achieve approximately $10^{7}$ faster than the numerical method, which completely satisfies the requirement of the QKD system in practice. In contrast, the numerical method takes several minutes to calculate a point of key rate, which is intolerable since many free-space sessions, such as satellite-ground or handheld QKD might have a window of only minutes. While collecting enough data based on the numerical method to train the model consumes a large amount of computing power, we can consider these large computations offline. Once we obtain the trained neural network, it can be deployed on a certain device to infer key rates online in milliseconds by giving new inputs from the experiment. Ref.Wang and Lo (2019) demonstrated that a neural network method for parameter optimization of QKD can be deployed on various mobile low-power systems, which brings advantages of more power efficiency and low latency. We can also forecast that our neural network method combined with Bayesian optimization will play an essential role in free-space QKD scenarios such as handheld Mélen et al. (2017), drone-based Hill et al. (2017) or satellite-ground QKD Liao et al. (2017). Several works have focused on machine learning for optimal parameters in QKD Wang and Lo (2019); Liu et al. (2018); Lu et al. (2019); Ding et al. (2020). However, our work predicts secure key rates directly by automatically designed neural networks, which goes further than our previous work Zhou et al. (2021).

Based on our model, there are several directions worthy of investigation for future work. Up to now, we have only covered computing the asymptotic key rates. However, finite-size effects are practical issues considered in discrete-modulated CV-QKD Leverrier et al. (2010). Note that a recent work has analyzed the security and performance of discrete-modulated CV-QKD under a finite-size scenario Almeida et al. (2021), which inspires us to improve our model. To address these issues, we also consider applying our model to other protocols in future work. Moreover, the issue of post-processing (notably the error correction part) still limits the overall time acceleration for a discrete-modulated continuous-variable QKD system. Note that the error correction involving binary or quaternary error-correcting codes is less complex compared with the situation of Gaussian modulation. Therefore, we also consider developing an effective error-correction protocol for CV QKD with discrete modulation using machine learning techniques in the future.

We gratefully acknowledge the support from the Natural Science Foundation of Jiangsu Province (No. BK20211145), the Fundamental Research Funds for the Central Universities (No. 020414380182), the Key Research and Development Program of Nanjing Jiangbei New Aera (No. ZDYD20210101), the Program for Innovative Talents and Entrepreneurs in Jiangsu(No. JSSCRC2021484), and the Key-Area Research and Development Program of Guangdong Province (No. 2020B0303040001). The authors would like to thank the High Performance Computing Center of Nanjing University for the numerical calculations.