Enhancing quantum state tomography via resource-efficient
  attention-based neural networks
Resource-efficient quantum state tomography is one of the key ingredients of
future quantum technologies. In this work, we propose a new tomography protocol
combining standard quantum state reconstruction methods with an attention-based
neural network architecture. We show how the proposed protocol is able to
improve the averaged fidelity reconstruction over linear inversion and
maximum-likelihood estimation in the finite-statistics regime, reducing at
least by an order of magnitude the amount of necessary training data. We
demonstrate the potential use of our protocol in physically relevant scenarios,
in particular, to certify metrological resources in the form of many-body
entanglement generated during the spin squeezing protocols. This could be
implemented with the current quantum simulator platforms, such as trapped ions,
and ultra-cold atoms in optical lattices.
Adriano Macarone Palmieri, Guillem Müller-Rigat, Anubhav Kumar Srivastava, Maciej Lewenstein, Grzegorz Rajchel-Mieldzioć, Marcin Płodzień
# Enhancing quantum state tomography via resource-efficient attention-based neural networks

###### Abstract

Resource-efficient quantum state tomography is one of the key ingredients of future quantum technologies. In this work, we propose a new tomography protocol combining standard quantum state reconstruction methods with an attention-based neural network architecture. We show how the proposed protocol is able to improve the averaged fidelity reconstruction over linear inversion and maximum-likelihood estimation in the finite-statistics regime, reducing at least by an order of magnitude the amount of necessary training data. We demonstrate the potential use of our protocol in physically relevant scenarios, in particular, to certify metrological resources in the form of many-body entanglement generated during the spin squeezing protocols. This could be implemented with the current quantum simulator platforms, such as trapped ions, and ultra-cold atoms in optical lattices.

+
Footnote †: These authors contributed equally.

+
Footnote †: These authors contributed equally.

+
Footnote †: These authors contributed equally.

## I Introduction

Modern quantum technologies are fuelled by resources such as coherence, quantum entanglement, or Bell nonlocality. Thus, a necessary step to assess the advantage they may provide is the certification of the above features [1; 2; 3; 4; 5; 6; 7; 8; 9; 10]. The resource content of a preparation is revealed from the statistics (e.g. correlations) the device is able to generate. Within the quantum formalism, such data is encoded in the density matrix, which can only be reconstructed based on finite information from experimentally available probes - a process known as quantum state tomography (QST) [11; 12; 13; 14; 15; 16; 17; 18; 19; 20; 21]. Hence, QST is a prerequisite to any verification task. On the other hand, the second quantum revolution brought new experimental techniques to generate and control massively correlated states in many-body systems [22; 23; 24; 25; 26], challenging established QST protocols. Both reasons elicited an extremely active field of research that throughout the years offered a plethora of algorithmic techniques [27; 28; 29; 30; 31; 32; 33; 34; 35; 36].

Over the recent years, machine learning (ML), artificial neural networks (NN), and deep learning (DL) entered the field of quantum technologies [37], offering many new solutions to quantum problems, also in the context of QST [38; 39; 40; 41; 42; 43]. The supervised DL approaches were already successfully tested on experimental data [44; 45; 46; 47], wherein a study of neural network quantum states beyond the shot-noise appeared only recently [48]. Nevertheless, neural network-based methods are not exempted from limitations, i.e. NN can provide non-physical outcomes [44; 46; 49], and are limited in ability to learn from a finite number of samples [50].

In this work, we address the abovementioned drawback by offering a computationally fast general protocol to back up any finite-statistic QST algorithms within a DL-supervised approach, greatly reducing the necessary number of measurements performed on the unknown quantum state. We treat the finite-statistic QST reconstructed density matrix as a noisy version of the target density matrix, and we employ a deep neural network as a denoising filter, which allows us to reconstruct the target density matrix. Furthermore, we connect the measured error loss function with quantum fidelity between the target and the reconstructed state. We show that the proposed protocol greatly reduces the required number of prepared measurements allowing for high-fidelity reconstruction \(\sim 98\%\) of both mixed and pure states. We also provide an illustration of the power of our method in, a physically useful, many-body spin-squeezing experimental protocol.

The paper is organized as follows: in Sec. II we introduce the main concepts behind the QST; in Sec. III we introduce the data generation protocol and neural network architecture, as well as define QST as a denoising task. Sec. IV is devoted to benchmarking our method against known approaches, and we test it on quantum states of physical interest. In Sec. V, we provide practical instructions to implement our algorithm in an experimental setting. We conclude in Sec. VI with several future research directions.

## II Preliminaries

Let us consider \(d\)-dimensional Hilbert space. A set of informationally complete (IC) measurement operators \(\hat{\mathbf{\pi}}=\{\hat{\pi}_{i}\}\), \(i=1,\ldots,d^{2}\), in principle allows to unequivocally reconstruct the underlying target quantum state \(\hat{\tau}\in\mathbb{C}^{d\times d}\) in the limit of an infinite number of ideal measurements [51; 52]. After infinitely many measurements,
one can infer the mean values:

\[p_{i}=\mathrm{Tr}[\hat{\tau}\hat{\pi}_{i}], \tag{1}\]

and construct a valid vector of probabilities \(\mathbf{p}=\{p_{i}\}\) for any proper state \(\hat{\tau}\in\mathcal{S}\), whereby \(\mathcal{S}\) we denote the set of \(d\)-dimensional quantum states, i.e. containing all unit-trace, positive semi-definite (PSD) \(d\times d\) Hermitian matrices. Moreover, \(\hat{\mathbf{\pi}}\) can be considered as a set of operators spanning the space of Hermitian matrices. In such a case, \(\mathbf{p}\) can be evaluated from multiple measurement settings (e.g., Pauli basis) and is generally no longer a probability distribution. In any case, there exists a one-to-one mapping \(Q\) from the mean values \(\mathbf{p}\) to the target density matrix \(\hat{\tau}\):

\[Q: \mathcal{F}_{\mathcal{S}}\longrightarrow\mathcal{S} \tag{2}\] \[\mathbf{p}\longmapsto Q[\mathbf{p}]=\hat{\tau},\]

where \(\mathcal{F}_{\mathcal{S}}\) is the space of accessible probability vectors. In particular, by inverting the Born's rule Eq. (1), elementary linear algebra allows us to describe the map \(Q\) as,

\[Q[\mathbf{p}]=\mathbf{p}^{T}\mathrm{Tr}[\hat{\mathbf{\pi}}\hat{\mathbf{ \pi}}^{T}]^{-1}\hat{\mathbf{\pi}}. \tag{3}\]

The inference of the mean values \(\mathbf{p}\) is only perfect in the limit of an infinite number of measurement shots, \(N\rightarrow\infty\).

In a realistic scenario, with a finite number of experimental runs \(N\), we have access to frequencies of relative occurrence \(\mathbf{f}=\{f_{i}:=n_{i}/N\}\), where \(n_{i}\) is the number of times the outcome \(i\) is observed. Such counts allow us to estimate \(\mathbf{p}\) within an unavoidable error dictated by the shot noise and of amplitude typically scaling as \(1/\sqrt{N}\)[53]. With only the frequencies \(\mathbf{f}\) available, we can use the mapping \(Q\) for an estimate \(\hat{\rho}\) of the target density matrix \(\hat{\tau}\), i.e.,

\[\hat{\rho}=Q[\mathbf{f}], \tag{4}\]

In the infinite number of trials \(N\rightarrow\infty\), \(f_{i}=p_{i}\) and \(\hat{\rho}=\hat{\tau}\). Yet, in the finite statistics regime, as considered in this work, the application of the mapping as defined in Eq. (3) to the frequency vector \(\mathbf{f}\) will generally lead to nonphysical results (i.e. \(\hat{\rho}\) not PSD). In such case, as an example of proper mapping \(Q\), we can consider different methods for standard tomography tasks, such as Linear Inversion (LI), or Maximum Likelihood Estimation (MLE), see Appendix A. As operators \(\hat{\mathbf{\pi}}\), we consider positive operator-valued measures (POVMs) and the more experimentally appealing Pauli basis (check Appendix B).

## III Methods

This section describes our density matrix reconstruction protocol, data generation, neural network (NN) training, and inference procedure. In Fig. 1, we show how such elements interact within the data flow. In the following paragraphs, we elaborate on the proposed protocol in detail.

The first step in our density matrix reconstruction protocol, called _pre-processing_, is a reconstruction of density matrix \(\hat{\rho}\) via finite-statistic QST with frequencies \(\mathbf{f}\) obtained from measurement prepared on target state \(\hat{\tau}\). Next, we feed forward the reconstructed density matrix \(\hat{\rho}\) through our neural network acting as a noise filter - we call this stage a _post-processing_. In order to enforce the positivity of the neural network output, we employ the so-called Cholesky decomposition to the density matrices, i.e. \(\hat{\rho}=C_{\rho}C_{\rho}^{\dagger}\), and \(\hat{\tau}=C_{\tau}C_{\tau}^{\dagger}\), where \(C_{\rho,\tau}\) are lower-triangular matrices. Such decomposition is unique, provided that \(\hat{\rho}\), \(\hat{\tau}\) is positive [127]. We treat the Cholesky matrix \(C_{\rho}\) obtained from finite-statistic QST protocol, as a _noisy_ version of the target Cholesky matrix \(C_{\tau}\) computed from \(\hat{\tau}\). The aim of the proposed neural network architecture is to learn a _denoising filter_ allowing reconstruction of the target \(C_{\tau}\) from the _noisy_ matrix \(C_{\rho}\) obtained via finite-statistic QST. Hence, we cast the neural network training process as a supervised denoising task.

Figure 1: Schematic representation of the data pipeline of our QST hybrid protocol. Panel (a) shows data acquisition from a generic experimental set-up, during which the frequencies \(\mathbf{f}\) are collected. Next, panel (b) presents standard density matrix reconstruction; in our work, we test the computationally cheap LI method together with the expensive MLE, in order to better analyse the network reconstruction behaviour and ability. Panel (c) depicts the matrix-to-matrix deep-learning strategy for Cholesky matrices reconstruction. The architecture herein considered is a combination of convolutional layers for input and output and a transformer model in between. Finally, we compare the reconstructed state \(\hat{\rho}\) with the target \(\hat{\tau}\).
### Training data generation

To construct the training data set, we first start with generating \(N_{\text{train}}\) Haar-random \(d\)-dimensional target density matrices, \(\{\hat{\tau}_{m}\}\), where \(m=1,\ldots,N_{\text{train}}\). Next, we simulate experimental measurement outcomes \(\mathbf{f}_{m}\), for each \(\hat{\tau}_{m}\), in one of the two ways:

1. _Directly_ : When the measurement operators \(\mathbf{\hat{\pi}}\) form an IC-POVM, we can take into account the noise by simply simulating the experiment and extracting the corresponding frequency vector \(\mathbf{f}_{m}=\{n_{i}/N\}_{m}\), where \(N\) is the total number of shots (i.i.d. trials) and the counts \(\{n_{i}\}_{m}\) are sampled from the multinomial distribution.
2. _Indirectly_ : As introduced in the preliminaries (Sec. II), if a generic basis is used \(\mathbf{\hat{\pi}}\), \(\mathbf{p}_{m}\) is no longer necessarily a probability distribution. This is the case with the Pauli basis (as defined in Appendix B), that we exploit in our examples. Then, we can add a similar amount of noise, obtaining \(\mathbf{f}_{m}=\mathbf{p}_{m}+\delta\mathbf{p}_{m}\), where \(\delta\mathbf{p}_{m}\) is sampled from the multi-normal distribution \(\mathcal{N}(\mathbf{0},\sim 1/(2\sqrt{N}))\) of mean zero and isotropic variance, saturating the shot noise limit.

Having prepared frequency vectors \(\{\mathbf{f}_{m}\}\), we apply QST via mapping \(Q\), Eq. (4), obtaining set of reconstructed density matrices \(\{\hat{\rho}_{m}\}\). We employ the most rudimentary and scalable method, i.e. the linear inversion QST, however, other QST methods can be utilized as well. Finally, we construct the training dataset as \(N_{\text{train}}\) pairs \(\big{\{}\vec{C}_{\rho},\vec{C}_{\tau}\big{\}}\), where we use \(\vec{C}\) to indicate the vectorization (flattening) of the Cholesky matrix \(C\) (see Appendix C for definition).

### Neural network training

The considered neural network, working as a denoising filter, is a non-linear function \(h_{\mathbf{\theta}}\) preparing matrix-to-matrix mapping, in its vectorized form, \(h_{\mathbf{\theta}}:\vec{C}_{\rho}\to\vec{C}_{\mathbf{\theta}}\), where \(\mathbf{\theta}\) incorporates all the variational parameters such as weights and biases to be optimized. The neural network training process relies on minimizing the cost function defined as a mean-squared error (MSE) of the network output with respect to the (vectorization of) target density matrix \(\hat{\tau}\), i.e.

\[\mathcal{L}^{\text{MSE}}(\mathbf{\theta})=\|\vec{C}_{\tau}-\vec{C}_{\mathbf{\theta}} \|^{2}, \tag{5}\]

via the presentation of \(K\) training samples \(\{\hat{\rho}_{l}\}_{l=1}^{K}\). The equivalence between MSE and Hilbert-Schmidt (HS) distance is discussed in detail in Appendix C, where we also demonstrate that the mean-squared error used in the cost function, Eq. (5), is a natural upper bound of the quantum fidelity. Hence, the choice of the cost function, Eq. (5), not only is the standard cost function for the neural network but also approximates the target state in a proper quantum metric. To make the model's optimization more efficient and avoid overfitting, we add a regularizing term resulting in the total cost function \(\mathcal{L}^{\text{MSE}}(\mathbf{\theta})+\text{Tr}[C_{\mathbf{\theta}}C_{\mathbf{\theta} }^{\dagger}]\) (chapter 7 of Ref. [54]).

The training process results in an optimal set of parameters of the neural network \(\mathbf{\bar{\theta}}=\text{arg min}_{\mathbf{\theta}}\mathcal{L}\), and trained neural network \(h_{\mathbf{\theta}}\), which allows for the reconstruction of the target density matrix \(\hat{\tau}\) via Cholesky matrix \(C_{\bar{\rho}}\)[128], i.e.

\[\hat{\bar{\rho}}=\frac{C_{\bar{\rho}}C_{\bar{\rho}}^{\dagger}}{\text{Tr} \Big{[}C_{\bar{\rho}}C_{\bar{\rho}}^{\dagger}\Big{]}}\simeq\hat{\tau}, \tag{6}\]

where \(C_{\bar{\rho}}\) is reshaped from \(\vec{C}_{\bar{\rho}}=h_{\mathbf{\theta}}(\vec{C}_{\rho})\).

### Neural network architecture

Our proposed architecture gets inspiration from the other recent models [55; 56], combining convolutional layers with a transformer layer implementing a self-attention mechanism [57; 58]. A convolutional neural network extracts important features from the input data, while a transformer block distils correlations between features via the self-attention mechanism. The self-attention mechanism utilizes the representation of the input data as nodes inside a graph [59] and aggregates relations between the nodes.

The architecture of considered neural network \(h_{\mathbf{\theta}}\) contains two convolutional layers \(h_{\text{cnn}}\) and transformer layer \(h_{\text{tr}}\) in between, i.e.:

\[h_{\mathbf{\theta}}(\vec{C}_{\rho})=\tanh[h_{\text{cnn}}\circ h_{\text{tr}}\circ \gamma(h_{\text{cnn}})](\vec{C}_{\rho}), \tag{7}\]

where \(\gamma(y)=1/2y(1+\text{Erf}(y)/\sqrt{2})\), \(y\in\mathbb{R}\), is the Gaussian Error Linear Unit (GELU) activation function [60] broadly used in the modern transformers architectures, and \(\tanh(y)\) is the hyperbolic tangent, acting element-wise of NN nodes.

The first layer \(h_{\text{cnn}}\) applies a set of \(K\) fixed-size trainable one-dimensional convolutional kernels to \(\vec{C}_{\rho}\) followed by non-linear activation function, i.e. \(\gamma(h_{\text{cnn}}(\vec{C}_{\rho}))\to\{\mathbf{F}_{\text{cnn}}^{1},\ldots,\mathbf{F }_{\text{cnn}}^{K}\}\). During the training process, convolutional kernels learn distinct features of the dataset, which are next feedforward to the transformer block \(h_{\text{tr}}\). The transformer block \(h_{\text{tr}}\) distills the correlations between the extracted features from kernels via the self-attention mechanism, providing a new set of vectors, i.e. \(h_{\text{tr}}(\mathbf{F}_{\text{cnn}}^{1},\ldots,\mathbf{F}_{\text{cnn}}^{K})\to\{\mathbf{F}_ {\text{tr}}^{1},\ldots,\mathbf{F}_{\text{tr}}^{K}\}\). The last convolutional layer \(h_{\text{cnn}}\) provides an output \(\vec{C}_{\mathbf{\theta}}\), \(\tanh(h_{\text{cnn}}(\mathbf{F}_{\text{tr}}^{1},\ldots,\mathbf{F}_{\text{tr}}^{K}))\to \vec{C}_{\mathbf{\theta}}\), where all filter outputs from the last layer are added. Finally, we reshape the output into the lower-triangular form and reconstruct the density matrix via Eq. (6).

The training data and the considered architecture allow interpretation of the trained NN as a conditional de
biaser (for details see Appendix D). The proposed protocol cannot improve the predictions of unbiased estimators, however, any estimator that outputs valid quantum states (e.g., LI, MLE) must be biased due to boundary effects. In the given framework, the task of the NN is to learn such skewness and drift the distribution towards the true mean.

## IV Results and discussion

Having introduced the features of our NN-based QST enhancer in the previous section, here, we demonstrate its advantages in scenarios of physical interest. To this aim, we consider two examples.

As the first example, we consider an idealized random quantum emitter (see e.g. Refs. [61; 62] for recent experimental proposals) that samples high-dimensional mixed states from the Hilbert-Schmidt distribution. After probing the system via single-setting square-root POVM, we are able to show the usefulness of our NN upon improving LI and MLE preprocessed states. We evaluate the generic performance of our solution and compare it with recent proposals of NN-based QST, Ref. [46].

In the second example, we focus on a specific class of muti-qubit pure states of special physical relevance, i.e. with metrological potential as quantified by the quantum Fisher information (QFI). Such states are generated via the famous one-axis twisting dynamics [63; 64]. Here, the system is measured using operators \(\mathbf{\hat{\pi}}\) via local symmetric, informationally complete, positive operator-valued measures (SIC-POVM) or experimentally relevant the Pauli operators.

In the following, we discuss the two abovementioned scenarios.

### Reconstructing high-dimensional random quantum states

**Scenario**.- As a first illustration, we consider a set of \(N_{\rm trial}\) random trial states \(\{\hat{\tau}_{j}\}\), \(j=1,\ldots,N_{\rm train}\) with Hilbert space dimension \(d=9\) sampled from the HS distribution (see Appendix E). The first task consists of assessing the average quality reconstruction over such an ensemble to benchmark the generic performance of our NN.

We prepare measurements on each trial state \(\hat{\tau}_{j}\) via informationally complete (IC) square-root POVMs as defined in Eq. (14), and obtain state reconstruction \(\hat{\rho}_{j}\) via two standard QST protocols, i.e. bare LI and MLE algorithms, as well as by our neural network enhanced protocols denoted as LI-NN, and MLE-NN, see Fig. 1. Finally, we evaluate the quality of the reconstruction as the average of the square of the Hilbert-Schmidt distance

\[D_{\rm HS}^{2}(\hat{\rho}_{j},\hat{\tau}_{j})={\rm Tr}[(\hat{\rho}_{j}-\hat{ \tau}_{j})^{2}]. \tag{8}\]

**Benchmarking**.- Fig. 2(a) presents averaged HS distance square \(\overline{D^{2}}_{\rm HS}\) as a function of the number of trials states \(N_{\rm trial}\), obtained with bare LI, MLE algorithms, and neural network enhanced LI (LI-NN), and neural network enhanced MLE (MLE-NN). The training dataset contains \(N_{\rm train}=1250\) HS distributed mixed states. Our neural network-enhanced protocol improves over the LI, and MLE, i.e. \(\overline{D^{2}}_{\rm HS}\) is lower for LI-NN, and MLE-NN compared to LI and MLE algorithms for the same \(N_{\rm trial}\). For fewer trials \(N_{\rm trial}<10^{3}\), the post-processing marginally improves over the states reconstructed only from MLE. For a larger number of trials, the lowest \(\overline{D^{2}}_{\rm HS}\) is obtained for MLE results, performing better than other considered tomographic methods. To enhance MLE in the regime of a few samples, we propose an alternative method by incorporating, as a depolarization channel, the statistical noise (see Appendix F).

Next, Fig. 2(b) presents comparison between our protocol with the state-of-the-art density matrix reconstruction via neural network, Ref. [46], where the authors report better performance than the MLE and LI algorithms having \(N_{\rm train}=8\cdot 10^{5}\) training samples, for states belonging to a Hilbert space of dimension \(d\geq 3\). Fig. 2(b) presents \(\overline{D^{2}}_{\rm HS}\) as a function of \(N_{\rm trial}/N_{\rm train}\), where \(N_{\rm train}=1250\) for MLE-NN, and \(N_{\rm train}=5000\) for LI-NN. The advantage of our protocol lies in the fact, that the amount of the necessary training data is significantly smaller compared to the training data size in Ref. [46], in order to obtain similar reconstruction qual

Figure 2: Evaluation of the QST reconstruction measured by the mean value of the Hilber-Schmidt distance square, \(\overline{D^{2}}_{\rm HS}\), for different QST protocols, averaged over \(N_{\rm test}=1000\) target states. Panel (a) shows mean \(D_{\rm HS}^{2}\) for four QST protocols: LI (green dots), MLE (red squares), NN-enhanced LI (blue diamonds), and NN-enhanced MLE (orange crosses). Panel (b) shows \(\overline{D^{2}}_{\rm HS}\) as a function of ratio \(N_{\rm trial}/N_{\rm train}\) for our LI-NN model (blue squares), MLE-NN model (orange crosses), and for the network model proposed in Ref. [46] (violet triangles). A lower number \(N_{\rm train}\) shifts lines to the left, indicating resource efficiency. Our proposed protocol achieves competitive averaged HS reconstruction for the size of training data of an order of magnitude smaller compared to the method proposed in Ref. [46] (note \(x\)-axis values increase to the left). During models’ training, we set \(N_{\rm train}=1250\) for MLE-NN protocol, and \(N_{\rm train}=5000\) for the LI-NN. Lines are to guide the eye, shadow areas represent one standard deviation.
ity, which is visible as a shift of the lines towards higher values of \(N_{\rm trial}/N_{\rm train}\).

### Certifying metrologically useful entanglement depth in many-body states

**Scenario.**- In the last experiment, we reconstruct a class of physically relevant multi-qubit pure states. Specifically, we consider a chain of \(L=4\) spins-\(1/2\) (Hilbert space of dimension \(d=16\)). The target quantum states are dynamically generated during the one-axis twisting (OAT) protocol [63; 64]

\[\ket{\Psi(t)}=e^{-itJ_{z}^{2}}\ket{+}^{\otimes L}\, \tag{9}\]

where \(\hat{J}_{z}\) is the collective spin operator along \(z\)-axis and \(\ket{+}^{\otimes L}=[(|\uparrow\rangle+|\downarrow\rangle)/\sqrt{2}]^{\otimes L}\) is the initial state prepared in a coherent spin state along \(x\)-axis (orthogonal to \(z\)). The OAT protocol generates spin-squeezed states useful for high-precision metrology, allowing to overcome the shot-noise limit [65; 66; 67; 68], as well as many-body entangled and the many-body Bell correlated states [69; 70; 71; 72; 73; 74; 75; 76; 77; 78; 79; 80; 81; 82]. OAT states have been extensively studied theoretically [83; 84; 85; 86; 87; 88; 89; 90; 91; 92], and can be realized with a variety of ultra-cold systems, utilizing atom-atom collisions [93; 94; 95; 96], and atom-light interactions [97; 98]. The recent theoretical proposals for the OAT simulation with ultra-cold atoms in optical lattices effectively simulate Hubbard and Heisenberg models [99; 100; 101; 102; 103; 104; 105; 106; 107].

**Preliminary results.**- For the first task, we generate our data for testing and training with SIC-POVM operators. For the test set, we select 100 OAT states in evenly spaced times \(t\in(0,\pi)\) and assess the average reconstruction achieved by our trained NN on a training dataset with \(N_{\rm train}=10,500\) Haar random pure states. We compare the values with the average score obtained for generic Haar distributed states, which is the set the model was trained on (see Appendix E).

The qualities of reconstructions are shown in Table 1. First, we verify that the NN is able to improve substantially the OAT states, even though no examples of such a class of states were given in the training phase, which relied only on Haar-random states. Moreover, the OAT-averaged reconstruction values exceed the Haar reconstruction ones. We conjecture that this stems from the bosonic symmetry exhibited by the OAT states. This symmetry introduces redundancies in the density matrix which might help the NN to detect errors produced by the statistical noise. Finally, let us highlight that the network also displays good robustness to noise. Indeed, when we feed the same network with states prepared for \(N_{\rm trial}=10^{3}\) trials, we increase the reconstruction fidelity moving to 87% from 67%.

**Inferring the quantum Fisher information (QFI).**- Finally, we evaluate the metrological usefulness of the reconstructed states as measured by the quantum Fisher information, \(F_{Q}[\hat{\rho},\hat{G}]\). The QFI is a nonlinear function of the state and quantifies the sensitivity upon rotations generated by \(\hat{G}\). For more details, we refer the reader to Appendix G.

Here, we use the collective spin component \(\hat{J}_{\bf v}\) as the generator \(\hat{G}=\hat{J}_{\bf v}\), where the orientation \(\mathbf{v}\in\mathbb{S}^{2}\) is chosen so that it delivers the maximal sensitivity. The QFI with respect to collective rotations can also be used to certify quantum entanglement [68], in particular, the entanglement depth \(k\), which is the minimal number of genuinely entangled particles that are necessary to describe the many-body state. If \(F_{Q}[\hat{\rho},\hat{J}_{\bf v}]>kL\), then, the quantum state \(\hat{\rho}\) has at least \(k+1\) entanglement depth [108; 109]. In particular, for states with depth \(k=1\) (i.e., separable), due to the lack of entanglement detected, the metrological power is at most the shot-noise limit [110]. This bound is saturated by coherent spin states - like our initial (\(t=0\)) state for OAT evolution, \(\ket{+}^{\otimes L}\).

In Fig. 3, we present the evolution of the QFI (normalized by the coherent limit, \(L=4\)) for the OAT target states (top solid blue lines). The top row corresponds to SIC-POVM measurement operators \(\mathbf{\hat{\pi}}\), while the bottom row corresponds to projective measurements with Pauli operators. In all experiments, we use the same neural network previously trained on random Haar states with frequencies \(\{\mathbf{f}_{m}\}\) obtained from \(N_{\rm trial}=10^{4}\) measurement trials. The LI algorithm by itself shows entanglement (QFI \(>L\)) for a few times \(t\). By enhancing the LI algorithm via our NN protocol, we surpass the three-body bound (QFI/\(L=3\)), thus revealing genuine 4-body entanglement, which is the highest depth possible in this system (as it is of size \(L=4\)). For instance, let us note that at time \(t=\pi/2\), the OAT dynamic generates the cat state, \(\ket{\Psi(t=\pi/2)}=(e^{-i\pi/4}|+\rangle^{\otimes 4}+e^{+i\pi/4}|-\rangle^{ \otimes 4})/\sqrt{2}\), which is genuinely \(L\)-body entangled, and so it is certi

\begin{table}
\begin{tabular}{|c||l|l|l|} \hline  & \(10^{5}\) trials & \(10^{4}\) trials & \(10^{3}\) trials \\ \hline LI-OAT & \(94.3\pm 0.4\%\) & \(86.5\pm 0.1\%\) & \(67.0\pm 2.0\%\) \\ \hline NN-OAT & \(98.6\pm 0.2\%\) & \(97.8\pm 0.5\%\) & \(87.6\pm 4.5\%\) \\ \hline NN-Haar & \(96.9\pm 3.0\%\) & \(94.2\pm 3.3\%\) & \(81.1\pm 4.1\%\) \\ \hline \end{tabular}
\end{table}
Table 1: Comparison of average fidelity and its standard deviation between the reconstructed and the target states of size \(d=16\) for various QST methods (rows), with varying size of measurement trials \(N_{\rm trial}=10^{5},10^{4},10^{3}\), as indicated by the consecutive columns. The first row presents the average fidelity reconstruction for linear inversion QST, averaged over OAT states, evenly sampled from \(t=0\) to \(t=\pi\). Employing our neural network presents an enhancement over the bare LI, as shown in the second row for the same target set. Finally, the third row also shows data for NN-enhanced LI, but averaged over general Haar-random states. All the initial Born values are calculated by noiseless SIC-POVM.
### Discussion

Currently, existing algorithms for QST tend to suffer from either of two main problems: unphysicality of results or biasedness. We tackle the problem of increasing the quality of state reconstruction by enhancing any biased QST estimator with an attention-based neural network working as a denoising filter, see Figure 1. Such an approach corrects both unphysicality and bias, at least partially. Our NN architecture is based on convolutional layers extracting crucial features from input data, and attention-based layers distilling correlations between extracted features. This choice of architecture is motivated by its improved generalization ability enabling a significant reduction of the necessary amount of training data, compared to NN-based approaches not utilizing an attention mechanism. From the examples provided in this section, we infer that our NN-enhanced protocol outperforms other methods of QST in the domain of a small number of measurement samples, see Figure 2. This is especially important due to applications in realistic scenarios, where the trade-off between accuracy and resource cost is crucial.

The results presented in this contribution prove that, with the small number of samples, the comparison of neural network enhancement to pure linear inversion protocol favors our implementation, as can be deduced by fidelity of reconstruction in Table 1. Although the NN was trained on Haar-random pure states, it achieves even better performance on its subset of measure zero, namely one-axis twisted states. We conjecture that this is due to their underlying symmetries, which allows to efficiently learn and correct the noise pattern.

Furthermore, the metrological usefulness of our method is visible through its certification of quantum Fisher information and entanglement depth, see Fig. 3. The bare QST setup, without our NN post-processing, is not able to show entanglement (QFI \(>L\)) at any finite time, as well as it never certifies the full genuine 4-body entanglement. Both of these problems are dealt with by the action of NN-enhancement.

## V Concrete experimental implementation for quantum state verification

To recapitulate this contribution, as a complement to Fig. 1 and our repository provided in Ref. [111], we summarize the practical implementation of the protocol introduced in this work.

1. _Scenario_: We consider a finite-dimensional quantum system prepared in a target state \(\hat{\tau}\). Here, we aim to verify the preparation of a quantum state \(\hat{\tau}\) via QST. To this end, we set a particular measurement basis \(\hat{\mathbf{\pi}}\) to probe the system.
2. _Experiment_: After a finite number of experimental runs, from the countings, we construct the frequency vector \(\mathbf{f}\).
3. _Preprocessed quantum state tomography_: From the frequency vector \(\mathbf{f}\) and the basis \(\hat{\mathbf{\pi}}\), we infer the first approximation of the state \(\hat{\rho}\) via the desired QST protocol (e.g., one of those introduced in Appendix A).
4. _Assessing pre-reconstruction_: We evaluate the quality of the reconstruction by e.g., computing \(D_{\mathrm{HS}}^{2}(\hat{\tau},\hat{\rho})\), quantum fidelity or any other meaningful quantum metric. To improve such a score, we resort to our NN solution to complete a denoising task. As with any deep-learning method, training is required.

Figure 3: Time evolution of the normalized QFI during the OAT protocol for \(L=4\) qubits system. Solid blue lines represent QFI calculated for target quantum states. The mean values of QFI calculated from tomographically reconstructed density matrices are denoted by green-dashed (reconstruction via LI), and red-dotted lines (reconstruction via neural network post-processed LI outputs). Shaded areas mark one standard deviation after averaging over 10 reconstructions. Panels (a) and (b) correspond to LI protocol with SIC-POVM data, whereas (c) and (d) denote LI reconstruction inferred from Pauli measurements. In the upper row, the left (right) column corresponds to \(N_{\mathrm{trial}}=10^{3}\) (\(10^{4}\)) trials; in the lower row, the left (right) column reproduces an LI initial fidelity reconstruction of \(\sim 74\%(\sim 86\%)\). The red lines represent the whole setup with NN post-processing of data from corresponding green lines, indicating improvement over the bare LI method. The NN advantage over the bare LI method can be characterized by entanglement depth certification, as shown by the horizontal lines denoting the entanglement depth bounds ranging from the separable limit (bottom line, bold), to the genuine \(L\)-body limit (top line). In particular, the presence of entanglement is witnessed by QFI \(>L\), as shown by the violation of the separable bound (bold horizontal line).
5. _Training neural network_: Different training strategies can be implemented: 1. Train over uniform ensembles (e.g., Haar, HS, Bures etc.) if \(\hat{\tau}\) is a typical state or we do not have information about it. If we know certain properties of the target state, we can take advantage of it (see next items). 2. Train over a subspace of states of interest. For example, if we reconstruct OAT states (Section IVb), we may train only on the permutation-invariant sector. 3. Train with experimental data. For example, if we have a quantum random source to characterize (Section IVa), experimental data can be used in the training set [44]. In such a case, our demonstrated reduction of the training set size translates also to a reduction of the experimental effort.
6. _Feeding the neural network_: We feed-forward the preprocessed state \(\hat{\rho}\) in our trained matrix-to-matrix NN to recover the enhanced quantum state \(\hat{\rho}\).
7. _Assessing the neural network_: We compute the updated reconstruction metric on the post-processed state \(D_{\text{HS}}^{2}(\hat{\tau},\hat{\bar{\rho}})\). Finally, we assess the usefulness of the NN by comparing how smaller such value is compared to the pre-processed score \(D_{\text{HS}}^{2}(\hat{\tau},\hat{\rho})\).

The strength of our proposed protocol lies in its broad applicability, as the choice of the basis \(\hat{\mathbf{\pi}}\) and QST pre-processing method is arbitrary.

## VI Conclusions

We proposed a novel deep learning protocol improving standard finite-statistic quantum state tomography methods, such as Linear Inversion and Maximum Likelihood Estimation. Our network, based on the attention mechanism and convolutional layers, greatly reduces the number of required measurements, serving as a denoising filter to the standard tomography output. The versatility of our approach comes from the fact that the measurement basis and reconstruction method have only an implicit impact as our central algorithm works directly with the density matrix. The proposed method reduces the number of necessary measurements on the target density matrix by at least an order of magnitude, compared to DL-QST protocols finite-statistic methods.

We verified that our proposed method is able to improve over LI and MLE preprocessed states. Moreover, the inference stage was performed on out-of-distribution data, i.e., we tested our model on density matrices forming an infinitesimally small fraction of the training dataset, indicating the robustness of the proposed method. In particular, we test our model on 4-qubits spin-squeezed, and many-body Bell correlated states, generated during one-axis twisting protocol, with average fidelity \(\sim 98\%\). We demonstrated our NN improves the reconstruction of a class of physically relevant multi-qubit states, paving the way to use such novel methods in current quantum computers and quantum simulators based on spin arrays.

Our protocol can greatly advance other QST methods, for both arbitrary states as well as for special classes that scale reasonably with the number of particles, such as symmetric states [112, 113].

**Data and code availability**.- Data and code are available at Ref. [111].

###### Acknowledgements.

 We thank Leo Zambrano, Federico Bianchi, and Emilia Witkowska for the fruitful discussions. We acknowledge support from: ERC AdG NOQIA; Ministerio de Ciencia y Innovation Agencia Estatal de Investigaciones (PGC2018-097027-B-I00/10.13039/501100011033, CEX2019-000910-S/10.13039/501100011033, Plan National FIDEUA PID2019-106901GB-I00, FPI, QUANTERA MAQS PC2019-111828-2, QUANTERA DYNAMITE PCI2022-132919, Proyectos de I+D+1 "Retos Colaboracion" QUSPIN RTC2019-007196-7); MICIIN with funding from European Union NextGenerationEU(PRTR-C17.I1) and by Generalitat de Catalunya; Fundacio Cellex; Fundacio Mir-Puig; Generalitat de Catalunya (European Social Fund FEDER and CERCA program, AGAUR Grant No. 2021 SGR 01452, QuantumCAT U16-011424, co-funded by ERDF Operational Program of Catalonia 2014-2020); Barcelona Supercomputing Center MareNostrum (FI-2022-1-0042); EU (PASQuanS2.1, 101113690); EU Horizon 2020 FET-OPEN OPTologic (Grant No 899794); EU Horizon Europe Program (Grant Agreement 101080086 -- NeQST), National Science Centre, Poland (Symmonia Grant No. 2016/20/W/ST4/00314); ICFO Internal "QuantumGaudi" project; European Union's Horizon 2020 research and innovation program under the Marie-Sklodowska-Curie grant agreement No 101029393 (STREDCH). AKS acknowledges support from the European Union's Horizon 2020 Research and Innovation Programme under the Marie Sklodowska-Curie Grant Agreement No. 847517. M.P. acknowledges the support of the Polish National Agency for Academic Exchange, the Bekker programme no: PPN/BEK/2020/1/00317. Views and opinions expressed are, however, those of the author(s) only and do not necessarily reflect those of the European Union, European Commission, European Climate, Infrastructure and Environment Executive Agency (CINEA), or any other granting authority. Neither the European Union nor any granting authority can be held responsible for them.