All-optical Neural Network Quantum State Tomography

Quantum state tomography (QST) is a standard process of reconstructing quantum information of an unknown quantum state through measurements on its copies. QST is used to verify state preparation, exam state properties such as correlations, and calibrate experimental systems. It is a crucial part for almost all aspects of experimental quantum information processing, including quantum computing, quantum metrology, and quantum communication Bouchard et al. (2019); D’Ariano et al. (2003); Leonhardt (1995); Thew et al. (2002); Lvovsky and Raymer (2009); Rambach et al. (2021).

As an analog of “imaging” technique in quantum settings, QST is born to be a data science problem. Given limited copies of an unknown state $\rho$, we can extract its information via QST. QST is essentially an inverse problem, and such information recovering tasks are well suited to machine leaning. Quantum learning theory indicates that $\Theta\left(2^{2n}/\varepsilon^{2}\right)$ copies of $\rho$ are necessary and sufficient to learn $\rho$ up to trace distance $\epsilon$ O’Donnell and Wright (2016). Although the tremendous resource requirement makes full-state QST impractical for large-scale systems, several weaker quantum learning models (e.g., probably approximately correct learning  Aaronson (2007), online learning Aaronson et al. (2019), shadow tomography Aaronson (2019); Huang et al. (2020)) can exponentially reduce the computational resource for learning some 2-outcome measurement expectation values or “shadows”.

Artificial neural network (NN), as a powerful algorithm in machine learning to fit a specific function, has been widely used for solving quantum information problems, such as quantum optimal control Niu et al. (2019); An et al. (2020), quantum maximum entropy estimation Cao et al. (2021), Hamiltonian reconstruction Cao et al. (2020), etc. NNs have aalso been widely applied for QST applications, such as recovering the information of local-Hamiltonian ground states from local measurements Xin et al. (2019) efficiently, performing tomography on highly entangled state with large system size Torlai et al. (2018), mitigating the state-preparation-and-measurement (SPAM) errors in experiments Palmieri et al. (2020), and improving the state fidelity Quek et al. (2018); Ahmed et al. (2020a). Generative models with neural networks can also perform QST with dramatically lower cost Carrasquilla et al. (2019); Ahmed et al. (2020b).

In this work, we demonstrate neural network QST with an all-optical neural network (AONN). Several optical implementations for realizing full connected neural network hardware have been proposed and demonstrated recentlyWetzstein et al. (2020); Shastri et al. (2021); Woods and Naughton (2012); Shen et al. (2017a); Lin et al. (2018); Zuo et al. (2019). Optical computing takes advantages of the bosonic wave nature of light: superposition and interference give rise to its intrinsic parallel computing ability. Meanwhile light is the fastest information carrier in nature. AONN is promising for next-generation artificial intelligence hardware which provides high energy efficiency, low crosstalk, light-speed processing, and massive parallelism. As compared to the electronic version, AONNs are ideal to deal with visual signals and information which are naturally generated and coded in light, such as image recognition and vehicular automation. However, most AONN demonstrations are still restricted to linear computation only due to the lack of suitable nonlinearity at low light level for large amount of optical neurons Woods and Naughton (2012); Shen et al. (2017a); Lin et al. (2018) . Without nonlinear activation functions, an AONN is always equivalent to a single-layer structure which cannot be applied for “real” deep machine learning. This problem has not been solved until most recently optical nonlinearity based on electromagnetically induced transparency (EIT) Zuo et al. (2019, 2021), phase-change materials Feldmann et al. (2019), and saturated absorption Guo et al. (2021); Ryou et al. (2021) was implemented to realize artificial optical neurons for AONNs.

Figure  1 illustrates a general scheme of AONN QST. Firstly we collect the training data set from a known quantum state {$|\psi_{j}\rangle$} and the corresponding local measurements {$M_{j}$}. Secondly we train neural network under supervised learning with some nonlinear activation functions in its hidden neurons to obtain the optimal network parameters. Third We take the trained network parameters to configure the AONN and perform some fine adjustment to optimize the hardware performance. At last, we feed measurement data sets to the trained AONN to reconstruct unknown quantum states. To validate this scheme, in the following sections we start from a general discussion of QST with the computer simulated NN and then describe our AONN experimental approach.

We here consider a general $n$-qubit space with Pauli operators (removed the all identity term) defined as

where $\mathcal{P}=\{\sigma_{0}=I,\sigma_{1}=X,\sigma_{2}=Y,\sigma_{3}=Z\}$. Every term in $P$ is specified by its index $(i_{1},i_{2},\cdots,i_{n})$. Measuring every element in $P$ performs a QST for any $n$-qubit quantum state $\rho$. For instance, when $n=1$, we need to measure all three Paulis $X,Y,Z$ for QST. Clearly, the cardinality of $P$ grows exponentially with $n$. When $\rho$ is a pure state, one may use techniques to reduce the number of measurements for $n>1$. Compressed sensing is one the standard techniques to recover low-rank quantum states from randomly sampled Pauli operators Gross et al. (2010); Flammia et al. (2012).

When $\rho$ is a pure state, it can be written as a ket

where {$\ket{\phi_{k}}$} are the computational basis, and the amplitudes $a_{k}\in\mathbb{C}$ are normalized (i.e. $\sum_{k=1}^{2^{n}}(a_{k,\text{r}}^{2}+a_{k,\text{im}}^{2})=1$, where $a_{k,\text{r}}\in\mathbb{R}$ and $a_{k,\text{im}}\in\mathbb{R}$ are the real and imaginary parts of $a_{k}$ respectively). The measurement expectation values $P$ are $\vec{c}=\operatorname{tr}(\rho\cdot P)=(\operatorname{tr}(\rho P_{1}),\operatorname{tr}(\rho P_{2}),\cdots,\operatorname{tr}(\rho P_{4^{n}-1}))$. For a single qubit pure state $\alpha|0\rangle+\beta|1\rangle$, its density matrix can be expressed as

where $\vec{\sigma}=\{X,Y,Z\}$, and $\vec{c}=(\operatorname{tr}(\rho X),\operatorname{tr}(\rho Y),\operatorname{tr}(\rho Z))\equiv(\langle X\rangle,\langle Y\rangle,\langle Z\rangle)$.

For $n>1$, we consider compressed sensing to reduce the number of measured operators. To perform compressed sensing, one needs to randomly sample a set $P^{m}=\{P_{1},\cdots,P_{m}\}$ of $m$ Pauli operators from $P$, then use $\vec{c}=\operatorname{tr}(\rho\cdot P^{m})=(\operatorname{tr}(\rho P_{1}),\operatorname{tr}(\rho P_{2}),\cdots,\operatorname{tr}(\rho P_{m}))$ to recover the unknown state $\rho$, more precisely, the parameters of $\rho$. This can be took as a regression problem to estimate the function between $\vec{c}$ and parameters of $\rho$ (e.g. $a_{k,\text{r}}$ and $a_{k,\text{im}}$).

NNs are excellent tools for solving regression problems. In the compressed sensing inspired QST, the expectation values $\vec{c}$ from random-sampled $P^{m}$ are inputs to the network, and ($a_{k,\text{r}}$, $a_{k,\text{im}}$) are the outputs. Note that there are different types of NNs with various structures and training procedures. According to the law of parsimony (Occam’s razor), we use the simplest type of NNs in this letter – fully-connected, feed-forward NNs. That is, the neurons between nearest layers are fully connected with one another and the information only pass forward while training. The supervised training process is to compare the ideal outputs $(a_{k,\text{r}},a_{k,\text{im}})$ with current NN outputs, and update parameters embedded in the NN to minimize their difference.

We numerically trained computer-based NNs nonlinear activation functions for 1-qubit, 2-qubit and 3-qubit QST. For the 1-qubit system, the number of sampled operators $m\in[1,2,3]$; for the 2-qubit system, the number of sampled operators $m\in[6,8,10,12]$; for the 3-qubit system, $m\in[20,25,30,35,40]$. Plainly, $m$ equals the number of input neurons, and $n$ decides the number of output neurons. For each $m$, three sets of Pauli operators have been sampled and tested. Figure 2 plots the average fidelities (green bars) of both cases as functions of the number of randomly sampled Paulis. For the single qubit system, the fidelity reaches 99.99% with 3 paulis. For the 2-qubit system, the fidelity reaches 99.9% with 10 randomly sampled paulis [Fig. 2(a)]. For the 3-qubit system, a fidelity of higher than 99.9% requires more than 35 randomly sampled Paulis [Fig. 2(b)]. Details of training can be found in Appendix A.

Theoretically, a pure state $\rho$ is Uniquely Determined among All states (UDA) of a set of operators $F$ means that there is no other state, pure or mix, has the same expectation values while measuring $F$ Chen et al. (2013). In Ref. Ma et al. (2016), authors discovered two sets of Pauli operators, $P_{\text{2-UDA}}$ and $P_{\text{3-UDA}}$, that are UDA for all 2-qubit and 3-qubit pure states respectively. (See Appendix B for the particular sets $P_{\text{2-UDA}}$ and $P_{\text{3-UDA}}$.) Namely, they are special cases of Pauli operator sets that the map between expectation values and the measured state $\rho$ is bijective. Similarly, we apply NNs for these two sets of UDA operators and obtain the prediction fidelities of 99.9% for the 2-qubit case and 99.3% for the 3-qubit case (red triangles in Figure 2).

We remark that our UDA scheme is not readily scalable for larger systems, however there exist protocals with better scalability, e.g., compressed sensing Gross et al. (2010), shadow tomography Aaronson (2019); Huang et al. (2020), where NN can be also be naturally used to enhance the protocols. Also, our NN-based scheme can be adapted to quantum tomography in the optical system, by taking into account of physical constrains, which will be discussed in detailed in the next section.

In this first proof-of-principle experimental demonstration, we implement the single qubit space with light polarizations, i.e., horizontal polarization $|H\rangle=|0\rangle$ and vertical polarization $|V\rangle=|1\rangle$. Instead of making a full QST, here we focus our task to determine the phase parameter of a pure state $|\psi\rangle=\frac{1}{\sqrt{2}}(|H\rangle+e^{i\theta}|V\rangle)$. The experimental AONN-QST setup is displayed in Fig.  3. In conventional QST, an arbitrary polarization state can be reconstructed by measuring the expectation values of the three Pauli operators. Figure 3(a) illustrates such an optical measurement setup. A laser beam passes through a polarization beam splitter (PBS₁) and becomes horizontally polarized ($|H\rangle$). The target state $|\psi\rangle=\frac{1}{\sqrt{2}}(|H\rangle+e^{i\theta}|V\rangle)$ is prepared by passing this horizontally polarized light through a half-wave plate (HWP₁) and a quarter-wave plate (QWP₁). The fast axis of the HWP₁ is aligned with an angle $\pi/4-\theta/2$ to the horizontal direction. The fast axis of the QWP₁ is aligned with an angle $\pi/4$ to the horizontal direction. $\langle X\rangle$, $\langle Y\rangle$ and $\langle Z\rangle$ are obtained by sending the light polarization qubit state to the measurement units II, III and IV shown in Fig. 3(a). To determine $\langle Z\rangle$, we send the polarization qubit directly to PBS₂ which projects $|H\rangle$ and $|V\rangle$ into two photodetectors in the measurement unit III. The normalized differential output from these two photodetectors gives the value $\langle Z\rangle$. The same setup can also be used to determine $\langle X\rangle$ or $\langle Y\rangle$ by placing HWP₂ or QWP₂ before PBS₂ as shown in II or IV, respectively (see Appendix C for details).

We obtain a data set {$M_{i}$}=$\{|\phi_{i}\rangle:1-\langle X\rangle_{i},1-\langle Y\rangle_{i},1-\langle Z\rangle_{i}\}$ by varying the phase $\theta\in[0,\pi/2]$ in the qubit state $|\psi\rangle=\frac{1}{\sqrt{2}}(|H\rangle+e^{i\theta}|V\rangle)$ and use them to train our AONN in Fig. 3(b). The AONN comprises an input layer of 3 neurons, a hidden layer of 20 neurons and a single-neuron output layer Zuo et al. (2019, 2021). Figure 3(b) shows the optical layout of the AONN and its network structure diagram is displayed in Fig. 3(c). The three coupling laser beams in the optical input layer are generated by a spatial light modulator (SLM1) [Fig. 3(b)], lenses L2 and L3, and an aperture, as shown in unit I of Fig.3(b). The SLM1 is divided into 3 parts and each part is encoded with sine phase pattern $m\pi\sin(\frac{2\pi}{T_{mi}}i+\frac{2\pi}{T_{mj}}j)$, where $m$ is the modulation depth, $T_{mi}$ and $T_{mj}$ are the period of modulation along $x$ and $y$ directions and $i$ and $j$ are the pixel number along $x$ and $y$ directions. The sine phase encoded on SLM1 modulate the beams into separated beams at focal plane of lens L2 and the aperture behaves as a filter to keep the zero-order beam, whose intensity is determined by the modulation depth $m$. Thus, the three beams with designed intensity in generated and collimated propagate to the SLM2 through lens L3. These weighted beams, as the input vector, are incident on SLM2 which diffracts each beam into 20 directions with designed weight (See Appendix D for the algorithm to calculate the pattern encoded on SLM2). A Fourier lens L4 performs linear summation for the beams diffracted into the same direction and forms 20 spots on its front focal plane. Thus, the combination of SLM2 and L4 completes the first linear operation W₁ and generates the input to the hidden layer. We then image these 20 spots with lenses L5 and L6 to laser-cooled ⁸⁵Rb atoms in a two-dimensional (2D) magneto-optical trap (MOT) Metcalf and van der Straten (2003); Zhang et al. (2012), where these 20-spot coupling beam pattern spatially modulate the transparency of the atomic medium through electromagnetically induced transparency (EIT) Harris (1997); Fleischhauer et al. (2005). Another relatively weak collimated probe beam counter-propagates through the MOT and its spatial transmission is nonlinearly controlled by the 20-spot coupling beam pattern. Here the nonlinear optical activation functions are realized with EIT in cold atoms. The equation of nonlinear activation functions follows Equation 7. The image of the probe beam transmission pattern by lens L6 and L8 becomes the output of the 20 hidden neurons. SLM3 and Fourier lens L9 perform the second linear matrix operation W₂ and the output is recorded by a camera. The technical details of our AONN is described in ref. Zuo et al. (2019, 2021).

In this work because we encode data into light energy, the AONN can only handle positive values: Input, output, linear matrix elements, and input/output of nonlinear activation functions are all positive values Zuo et al. (2019, 2021). Meanwhile, the EIT optical nonlinear activation functions are increasing and convex. Therefore the AONN is only able to perform regression task on increasing and convex functions. To match the AONN constrains, we performing a transform to the input variable, e.g., $\langle X\rangle$ to 1-$\langle X\rangle$, so that all input values to the AONN nodes are positive. We add these conditions to NN to simulate the AONN performance. We find that this specific AONN fails to describe the whole range of nonmonotonic functions. For the first proof-of-principle experimental demonstration, we will only apply the AONN for single-qubit QST with phase $\theta$ within $[0,\pi/2]$. It is surprising that such a positive-valued AONN is still able to perform some types of QST.

To train the AONN, we prepared the training data set {$M_{i}$} from 23 phase values from a uniform distribution $\theta_{j}\sim U(0,\pi/2)$, corresponding to the optical polarization states $\{\rho_{j}=\mathcal{N}(|\phi_{j}\rangle\langle\phi_{j}|)\}$. Here $\mathcal{N}$ is the noise channel in experiments, and measure the Pauli expectation values $\langle X\rangle$, $\langle Y\rangle$, $\langle Z\rangle$. In a similar way, we prepare a test set with 32 independent data samples.

In addition to optical quantum states, we sample data from the IBM quantum computer ibmq_ourense ibm (2020), and implement the same AONN training for comparison. The quantum circuit to prepare $|\psi\rangle=(|H\rangle+e^{i\theta}|V\rangle)/\sqrt{2}$ is the initial state $|H\rangle$ going through a Hadamard gate and then going through a RZ rotation gate. On ibmq_ourense, we uniformly sample 158 data points as training set, 50 data points as test set. Both experimental optical quantum state tomography and IBMQ tomography data are used to train two NNs. Details of training AONN can be found in Appendix E.

Figure 4 shows the AONN state construction results using neural network models trained by the AONN QST training set and the IBM quantum computer training set separately. The theoretical value is calculated from $\langle X\rangle$ directly. With the AONN system set up to the training results, we sent a set of input vector to the system. The example of real and imaginary part of density matrix are shown in Fig. 4(a). The experimentally measured state example is predicted by AONN QST training model. The example input vector for AONN model is $(\langle X\rangle,\langle Y\rangle,\langle Z\rangle)=(0.440,0.898,0)$ and the AONN experimental predicted state is $\theta=1.195$ and $\rho=\begin{pmatrix}0.5&0.1852-0.46457i\\ 0.1852+0.46457i&0.5\end{pmatrix}$ which is close to the theoretical value $\theta=1.1152$ and neural network predicted value $\theta=1.1532$. The state is also marked with yellow triangle in Fig. 4 (b1). The experimental results are shown in Fig. 4(b). The theoretical value, NN predicted value and experimental AONN predicted value collapse together for both optics data training (Fig. 4 (b1) shows) and IBMQ data training (Fig. 4 (b2) shows). The theoretical value, neural network prediction value and AONN predicted value are consistent in both cases. The results suggest that our positive-valued AONN with EIT nonlinear activation functions is capable for qubit QST.

While most demonstrations of optical neural networks took classification tasks to verify their feasibilityLin et al. (2018); Feldmann et al. (2019); Shen et al. (2017b), we performed the first AONN-QST. To accomplish regression tasks, nonlinear function is essential as long as the relation between input vector and output vector cannot be expressed linearly. The tunable EIT nonlinear optical activation functions in our AONN offer opportunities for performing regression tasks with certain functions. Although our AONN has some certain limitation that the linear operation matrix elements are all positive valued, it has potential to solve some QST problems as we have demonstrated here for single qubit tomography. We believe the next generation of complex-valued AONNs with encoding data in both light amplitude and phase will be more powerful. The future development of complex-valued AONNs may become a more powerful tool for full QST in a much larger $n$-qubit space.

Optical quantum network Kimble (2008) has been brought to the fore by the reduced decoherence and high speed of photons. Recently, apart from generating optical quantum states Gu et al. (2019) and optical quantum communication over long distance Yu et al. (2020), multiple state-of-the-art experiments on optical quantum interfaces to store Li et al. (2020) and distribute entanglements Choi et al. (2010); Pu et al. (2018) have been exhibited. Among all of these, QST is essential for not only characterizing generation and preservation of quantum states but also has the potential to verifying the entanglement distributed across the whole network. We believe that our all-optical setup of integrated AONN-QST will shed light on replenishing the all-optical quantum network with one more brick. We acknowledge the use of IBM Quantum services for this work.