Classification and reconstruction of optical quantum states with deep neural networks

We apply deep-neural-network-based techniques to quantum state classification and reconstruction. Our methods demonstrate high classification accuracies and reconstruction fidelities, even in the presence of noise and with little data. Using optical quantum states as examples, we first demonstrate how convolutional neural networks (CNNs) can successfully classify several types of states distorted by, e.g., additive Gaussian noise or photon loss. We further show that a CNN trained on noisy inputs can learn to identify the most important regions in the data, which potentially can reduce the cost of tomography by guiding adaptive data collection. Secondly, we demonstrate reconstruction of quantum-state density matrices using neural networks that incorporate quantum-physics knowledge. The knowledge is implemented as custom neural-network layers that convert outputs from standard feed-forward neural networks to valid descriptions of quantum states. Any standard feed-forward neural-network architecture can be adapted for quantum state tomography (QST) with our method. We present further demonstrations of our proposed QST technique with conditional generative adversarial networks (QST-CGAN) [Ahmed et al., Phys. Rev. Lett. 127, 140502 (2021)]. We motivate our choice of a learnable loss function within an adversarial framework by demonstrating that the QST-CGAN outperforms, across a range of scenarios, generative networks trained with standard loss functions. For pure states with additive or convolutional Gaussian noise, the QST-CGAN is able to adapt to the noise and reconstruct the underlying state. The QST-CGAN reconstructs states using up to two orders of magnitude fewer iterative steps than iterative and accelerated projected-gradient-based maximum-likelihood estimation (MLE) methods. We also demonstrate that the QST-CGAN can reconstruct both pure and mixed states from two orders of magnitude fewer randomly chosen data points than these MLE methods. Our paper opens possibilities to use state-of-the-art deep-learning methods for quantum state classification and reconstruction under various types of noise.

Figure 1

An outline of the topics discussed in this work. (a) The data required for quantum state tomography consists of the frequencies of measurement outcomes from observables represented as Hermitian operators. The aim is to reconstruct a description of the quantum state—usually a density matrix or the wave function. (b) Measurements of quasi-probability distributions (Wigner or Husimi $Q$) reveal interesting visual features in the data. Similarly, the histograms of measurement statistics, e.g., the photon-number distribution, can have patterns. Such features and patterns can be used for classification or reconstruction. (c) Several types of noise can corrupt the state or the data. Some types of noise, e.g., white noise, can be reduced by more data collection. Other types of noise, e.g., state-preparation-and-measurement (SPAM) noise, are more difficult to handle. (d) Classification tasks attempt to assign a label to data, classifying it according to its properties, e.g., if it is a Schrödinger-cat state, has Wigner negativity, or is an entangled quantum state. (e) Neural networks can be trained for classification of states or their properties. (f) Once it has been trained, analyzing how a neural network determines the class of a state can help to focus on the most important features in the data. This can be leveraged for adaptive data collection. (g) Reconstruction of quantum states connects to generative modeling tasks, where the goal is to learn the underlying probability distribution of the data to sample new data from it. In quantum state reconstruction, we aim to learn an underlying model, usually in the form of a wave function or density matrix that can generate statistics for any measurement operator. (h) Neural-network methods can be used for estimating or approximating underlying probability densities explicitly using restricted Boltzmann machines (RBMs) or variational autoencoders (VAEs). Training RBMs is not straightforward due to sampling requirements. This is resolved in VAEs using a reparameterization trick that allows gradient-based backpropagation for training. (i) Generative adversarial networks (GANs) provide a density-estimation technique, where we do not explicitly define the density nor require the reparametrization trick for training. We combine ideas from VAEs and GANs to propose a new quantum state tomography technique with conditional GANs—the QST-CGAN. Our QST-CGAN method allows for explicit estimation of the density matrix and computation of measurement statistics using two custom layers—DensityMatrix and Expectation.

Figure 2

A $\mathtt{binomial}(S=2,N=5,\mu =0)$ state and data generated from a displace-and-measure calculation using QuTiP [150, 151] within a $200\times 200$ grid. (a) The photon occupation probabilities, i.e., the diagonal elements of the density matrix $\rho$. (b) A Hinton plot of $\rho$, where blue (red) denotes that the real part of the density-matrix element is positive (negative). The size and the shade of each square is determined by the absolute value of the density-matrix element. [(c),(d),(e)] The generalized $Q$ function, $Q_n^{\beta }$, for $n=0,1,2$. (f) The corresponding Wigner function computed using the different $Q_n^{\beta }$ as $(2/\pi ){\sum }_n{(-1)}^n{Q}_n^{\beta }$. Note that even when choosing a Hilbert-space cutoff of 100 for this demonstration, the corners in the Wigner-function plot have spurious nonzero values at large displacements $\beta \approx \pm 5\pm 5i$. To mitigate such effects, larger cutoffs are required for states that have a high photon number or we need to restrict the computation to smaller values of $\beta$. Other methods of computing the Wigner function from $\rho$ do not suffer such problems even with a cutoff of 16 for this specific example. QuTiP provides several such implementations and we use one of them, the numerically stable Clenshaw method, to compute Wigner functions in the rest of the paper.

Figure 3

Representative examples from each class of optical quantum states considered in this paper. In the top row, we plot the Wigner function for the states, using the same scaling as Fig. 2. In the bottom row, we show the values of the density-matrix elements for each state as Hinton plots similar to Fig. 2. We can see that the Wigner functions and density matrices have characteristic patterns that a neural network can learn and use for classification or reconstruction.

Figure 4

The effect of various types of noise on the measurement data (Husimi $Q$ functions) for the state in Fig. 2. We normalize the data to $[0,1]$ by dividing with the maximum value and will use this color scheme throughout the paper to represent such rescaled data. (a) Mixed states with $\zeta =0.5$ and density 0.8 for ${\rho }_{\text{random}}$. (b) Convolution with Gaussian noise with $n_{\mathrm{th}}=3$. (c) Photon loss with $\gamma \tau$ such that 50% of the average initial photons have been lost. (d) Affine transformation with rotation $\theta ={100}^{\text{o}}$, shear $\mathrm{\Omega }=5^{\text{o}}$, and $\mathrm{\Delta }x=\mathrm{\Delta }p=1.61$ (5 pixels). (e) Additive Gaussian noise with standard deviation $\sigma =0.2$. (f) Pepper noise setting 50% of the data points to zero.

Figure 5

Sketch of the Classifier network, which classifies optical quantum states from Husimi $Q$ data (input image on the left). The blocks represents convolution operations, where filters (exemplified by boxes connecting one layer to another) extract features from the image. We use six such convolutional layers in our architecture (we only show three here). The extracted features are fed to the first of three fully connected layers. The outputs of the last layer are converted to a classification label. For the parameters used, see Table 2.

Figure 6

Sketch of the Generator $G$ and Discriminator $D$ neural networks adapted for quantum state reconstruction. The two inputs to $G$ are the measurement statistics $\mathbf{d}$ and the observables $\left\{{\mathcal{O}}_i\right\}$. The input $\mathbf{d}$ is taken as a flattened vector, which is reshaped to a $16\times 16\times 2$ matrix after the first layer of $G$. Then, after successive transpose convolution operations, we obtain a $32\times 32\times 2$ matrix. This intermediate output is converted into a lower-triangular matrix with real elements on the diagonal to obtain a Cholesky decomposition form, $T_G$, that can yield a valid density-matrix representation ${\rho }_G$. The expectation values of the observables are then computed using the Born rule $\text{tr}\left({\mathcal{O}}_i{\rho }_G\right)$. In the last layer, any known source of noise is added to the outputs. The inputs to $D$ are a concatenation of $\mathbf{d}$ with either the generated data from $G$ or $\mathbf{d}$ itself. The output is interpreted as a similarity score between the inputs (the score is $\sim 1$ if they match, i.e., for inputs $\sim \mathbf{d}$). The weights of the two networks, ${\theta }_G$ and ${\theta }_D$ are updated alternatingly to minimize their respective loss functions. For the details of all the parameters, see Table 3 and Table 4.

Figure 7

Performance of the $\mathtt{Classifier}$. (a) Confusion matrix demonstrating the performance of the $\mathtt{Classifier}$ on a test dataset containing 8760 states ($\sim 1200$ different instances of each class). The prediction counts are normalized to show the true labels versus the predictions made by the $\mathtt{Classifier}$. (b) Husimi $Q$ functions for three examples of pure cat states that the $\mathtt{Classifier}$ does not classify correctly. For each state, the incorrectly assigned label is shown. The states are, from left to right, $\mathtt{cat}(\alpha =1,S=1,\mu =1)$, $\mathtt{cat}(\alpha =2,S=3,\mu =0)$, and $\mathtt{cat}(\alpha =1,S=3,\mu =0)$. For certain parameters, different states have a high overlap in fidelities and the measurement data, making classification challenging. The $\mathtt{Classifier}$ tries to find the best label according to relevant patterns in the data.

Figure 8

$\mathtt{Classifier}$ performance for cat states with photon loss. (a) Softmax probabilities (solid lines) predicted by the $\mathtt{Classifier}$ for the labels of the seven classes. The shaded regions show one standard deviation from the mean. The dataset consists of 100 cat states ($\left|\alpha \right|\in [2,3],S=0,\mu =0$) with photon loss, quantified by the proportion of photons lost $\delta n/n_0$ starting from the initial mean photon number $n_0$. (b) Husimi $Q$ functions for one of the cat states in the dataset with, from left to right, $0\%$, $20\%$, and $100\%$ of photons lost with respective fidelities 0.76 and 0.19 for the states with photon loss. It is not straightforward to assert from just the Husimi-$Q$ data when a $\mathtt{cat}$ state stops being a “cat” as it still possesses $\mathtt{cat}$-like features (two coherent blobs) even after losing $20\%$ of the initial photons. (c) The photon-number occupation probabilities for the states in (b). Note that the occupation probability for the vacuum state is $\sim 1$ in the right panel, but we set the limits of the y axis to 0.5 for better distinguishability.

Figure 9

$\mathtt{Classifier}$ performance in the presence of additive Gaussian noise. (a) Accuracy (blue) and area under the ROC curve (orange) as a function of the level of additive Gaussian noise. The noise is added to the dataset used in Fig. 7. (b) Effect of additive Gaussian noise on a $\mathtt{cat}(\alpha =2,S=0)$ state. The distortion of the state starts to become significant between $\sigma =0.05$ and $\sigma =0.2$, which is when the accuracy of the Classifier starts to drop below 90%.

Figure 10

Using Grad-CAM to highlight the regions of the input data that the Classifier considers most important for predicting a label. (a) Input data for three states (from left to right: binomial, num, and gkp) with affine transformations and a constant additive Gaussian noise ($\sigma =0.2$) for all values of $\beta$ after normalizing the data to the interval $[0,1]$. (b) Heatmaps, normalized to the interval $[0,1]$, constructed with Grad-CAM from the data in (a), showing which parts of the data the Classifier focusses on. (c) The areas of the data (without the additive Gaussian noise) that appear in the focus when we only show the regions for which the Grad-CAM signals exceed 0.9.

Figure 11

The effect of the loss function on reconstruction of $\rho$ using the $\mathtt{Generator}$. We use the data from the Husimi $Q$ function of a $\mathtt{binomial}(S=2,N=4,\mu =0)$ state [inset in panel (d)] in a $32\times 32$ grid and repeat the reconstruction with random initializations of the network weights and starting estimate of $\rho$ for iMLE. In each iteration, the weights of the $\mathtt{Generator}$ or $\mathtt{Discriminator}$ networks are updated using a single step of the Adam optimizer. The learning-rate schedule and optimization hyperparameters are set to the same values (see Sec. 4b3) for all loss functions in order to achieve a fair comparison. However, further tuning of the parameters for each type of loss function could possibly give better results. (a) The reconstruction fidelity as a function of iterations for QST-CGAN with various weights of the L1 loss set by the ${\lambda }_{\mathrm{L}1}$ parameter [see Eq. (54)]. In each of a total of 30 runs, the weights of the $\mathtt{Generator}$ and $\mathtt{Discriminator}$ are randomly initialized. The solid lines show the mean and the shaded regions shows one standard deviation from the mean. (b) The performance of MLE methods on the same data. We repeat the reconstruction 30 times and show the mean fidelity (solid-blue line) and one standard deviation from the mean (shaded region) for iMLE. For the APG-MLE (dashed-black line), we use the default initialization scheme (“bootstrap”) from Ref. [119], which initializes the starting density matrix via conjugate gradients using a line search. The plot for APG-MLE shows the improvement of fidelity including the steps during the initialization scheme. There is no deviation from the mean for APG-MLE, since there is no explicit randomization involved in the reconstruction. [(c),(d),(e),(f)] Reconstruction fidelities using standard loss functions for the $\mathtt{Generator}$: cross-entropy [see Eq. (49)], KL divergence, L1, and L2 [see Eq. (50)]. We show all 30 runs for each loss function with the dashed line showing the mean.

Figure 12

The effect of the loss function on reconstruction of a $\mathtt{cat}(\alpha =2,S=0,\mu =0)$ state from Husimi-$Q$-function data [inset in panel (b)]. All hyperparameters, number of runs, and meanings of solid lines and the shaded regions in the plots are the same as for Fig. 11. (a) Performance of the QST-CGAN with various weights of the L1 loss. (b) Reconstruction fidelities for MLE methods on the same data. [(c)–(f)] Reconstruction performance with standard loss functions. The inset in (d) shows the Husimi $Q$ function of a $\mathtt{cat}$ state orthogonal to the one in the inset in (b), which was used to produce the data. The state in (d) is constructed by applying the photon annihilation operator $a$ to the original state in (b).

Figure 13

Reconstruction of a $\mathtt{binomial}(S=2,N=4,\mu =0)$ state in the presence of additive Gaussian noise. (a) The Husimi $Q$ function of the state after addition of Gaussian noise at each $\beta$. The random noise is drawn from a standard normal distribution with $\sigma =0.05$ and added after the data has been normalized to the range $[0,1]$. [(b)–(f)] Reconstructed Husimi $Q$ functions, without noise added by the GaussianNoise layer, using standard loss functions for the $\mathtt{Generator}$: L1, cross entropy, L2, and KL divergence, respectively. (d) Reconstructed Husimi $Q$ function using APG-MLE. [(g)–(i)] Reconstructed Husimi $Q$ functions using our QST-CGAN with three different weights of the L1 loss set by ${\lambda }_{\mathrm{L}1}$. (j) Photon-number occupation probabilities for the data without noise added. [(k)–(r)] Photon-number occupation probabilities extracted from the reconstructed density matrices corresponding to the Husimi $Q$ functions in (b), (c), (d), (e), (f), (g), (h), and (i), respectively. In all reconstructions using neural networks, the hyperparameters for learning were kept the same. For each method, including APG-MLE, the calculations were stopped after 10 000 iterations.

Figure 14

The effect of the loss function on reconstruction of the binomial state in Fig. 13 with 30 different realizations of the additive Gaussian noise. We show all runs for each loss function. The injected noise has the same variance $\sigma =0.05$ in each run, but is sampled anew for each run. [(a)–(c)] Reconstruction fidelity for the QST-CGAN with various weights of the L1 loss. (d) Reconstruction fidelity for MLE methods. [(e)–(h)] Reconstruction fidelity obtained by training the $\mathtt{Generator}$ using standard loss functions: cross entropy, KL divergence, L1, and L2. In all the neural-network-based reconstructions, we use the same hyperparameters for training in order to have a fair comparison.

Figure 15

Reconstruction in the presence of Gaussian convolution noise. We assume that the noisy amplification channel has a thermal noise with $n_{\text{th}}=5$ photons (see Sec. 3c2). This noise is convoluted with the data from the Husimi $Q$ function in (a) a $\mathtt{fock}(n=1)$ state and (b) a $\mathtt{binomial}(S=2,N=4)$ state. We had to use an $81\times 81$ grid for the data since we noticed that the convolution operation leads to effects such as loss of radial symmetry in (a) when using a grid of $32\times 32$. The image obtained by subtracting the background from the convoluted data shows the symmetries of the state. We show the reconstruction of the underlying states in by the QST-CGAN (${\lambda }_{\mathrm{L}1}=10$) where the $\mathtt{Generator}$ output is convoluted with the same background noise. We recover the underlying state from the DensityMatrix layer of the $\mathtt{Generator}$. Note that even though the convoluted outputs and subtracted counts match the data well in both cases, the fidelity between the underlying reconstructed state itself and the true underlying state is 1 for (a), but only 0.45 in (b).

Figure 16

Reconstruction of a mixture $\rho =0.8\mathtt{cat}+[0.2/(r-1)]{\sum }_{n=0}^{r-2}\mathtt{fock}\left(n\right)$ of $\mathtt{cat}(\alpha =2,S=0,\mu =0)$ (the state used in Fig. 12) and $\mathtt{fock}\left(n\right)$ states, where $r\ge 2$ denotes the rank. For $r=1$, the state is just $\mathtt{cat}(\alpha =2,S=0,\mu =0)$. The input data is the Husimi $Q$ function of $\rho$ measured in a $32\times 32$ grid. The solid lines show the mean and the shaded regions show one standard deviation from the mean for the QST-CGAN (red) and iMLE (blue) over 15 reconstructions for each of ranks $r=1,2,3,4$. The dashed black lines show the fidelities for states given by APG-MLE using the default “bootstrap” initialization. The APG-MLE method does not have any randomization and therefore does not have a standard deviation from the mean. In each repetition, we use the same data, but start from a different random initial state for iMLE and random weights for the QST-CGAN. We choose the weight ${\lambda }_{\mathrm{L}1}=1$ for the QST-CGAN and keep all other training hyperparameters the same as used for previous results and described in Sec. 4b3. The QST-CGAN runs are stopped when they have converged on a reconstructed state.

Figure 17

Reconstruction of a thermal state with mean photon number $n_{\mathrm{th}}=1$. The QST-CGAN method uses the weight ${\lambda }_{\mathrm{L}1}=1$; all other training hyperparameters are kept the same as for previous results and described in Sec. 4b3. (a) The fidelity of the reconstructed state as a function of the number of iterations for iMLE (blue) and QST-CGAN (red). [(b)–(d)] Reconstructed data compared to the data used for obtaining the underlying density matrix. We use the Husimi $Q$ function measured in a $32\times 32$ grid as the input. [(e)–(g)] Wigner function of the underlying thermal state compared to the Wigner functions obtained from the reconstructions. The reconstructed Wigner functions do not match the smooth nature of the Wigner function obtained from the underlying state.

Figure 18

Reconstruction with a QST-CGAN of a random mixed state of rank 11 from Husimi $Q$ and Wigner functions sampled on a $32\times 32$ grid. [(a),(b)] Data used for the reconstruction: the Husimi $Q$ and Wigner functions. (c) Hinton plot (see Fig. 3) of the underlying density matrix. [(d),(e)] Husimi $Q$ and Wigner functions for a $64\times 64$ grid computed from the underlying state to show finer features not present in the data that is fed to the neural network. [(f)–(h)] Reconstruction results using the Husimi-$Q$-function data in (a) as input for the QST-CGAN. [(i),(j),(k)] Reconstruction results using the Wigner-function data in (b) as input for the QST-CGAN.

Figure 19

Reconstruction of a rank-4 mixture of $\mathtt{cat}(\alpha =2,S=0,\mu =0)$ and $\mathtt{fock}\left(n\right)$ states (the mixture is constructed using the same formula as in Fig. 16) from input data consisting of 256 points of the Husimi $Q$ function (instead of the full $32\times 32=\text{1024}$ points used in preceding sections). (a) Reconstruction fidelity of the QST-CGAN (with ${\lambda }_{\mathrm{L}1}=1$; red) and iMLE (blue) where we plot all trajectories for 36 different reconstructions and show the mean with dashed lines. In each reconstruction, we randomly selected a set of $\beta$ values from the phase space and reconstruct the state where the QST-CGAN weights and the initial state for iMLE are reinitialized randomly. [(b)–(d)] Comparison between the data and the reconstructed Husimi $Q$ function given by iMLE and QST-CGAN for one selection of input data points (white points in the left panel). [(e)–(g)] Photon-number occupation probabilities from the (reconstructed) density matrices in (b).

Figure 20

Reconstruction of a mixture of $\mathtt{cat}(\alpha =2,S=0,\mu =0)$ and $\mathtt{fock}\left(n\right)$ states (the mixture is constructed using the same formula as in Fig. 16) with a reduced number of measurements. Reconstruction fidelities for ranks $r=2,3$, and 4 are shown for QST-CGAN (red) and iMLE (blue). The solid lines show the mean and the shaded regions show one standard deviation from the mean. The fidelity shown is the one reached after a certain convergence criterion set by a tolerance value. We choose a tolerance such that if the average fidelity in 100 iterations does not change by ${10}^{-5}$ over 5 steps (i.e., 500 iterations) we stop the reconstruction.