QuaDiM: A Conditional Diffusion Model For Quantum State Property Estimation
Authors: Yehui Tang
Mabiao Long
Junchi Yan
Venue: ICLR 2025 Poster (2025)
Keywords: quantum
property estimation
machine learning
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: Quantum state property estimation (QPE) is a fundamental challenge in quantum many-body problems in physics and chemistry, involving the prediction of characteristics such as correlation and entanglement entropy through statistical analysis of quantum measurement data. Recent advances in deep learning have provided powerful solutions, predominantly using auto-regressive models. These models generally assume an intrinsic ordering among qubits, aiming to approximate the classical probability distribution through sequential training. However, unlike natural language, the entanglement structure of qubits lacks an inherent ordering, hurting the motivation of such models. In this paper, we introduce a novel, non-autoregressive generative model called \textbf{\model}, designed for \underline{\textbf{Qua}}ntum state property estimation using \underline{\textbf{Di}}ffusion \underline{\textbf{M}}odels. \model progressively denoises Gaussian noise into the distribution corresponding to the quantum state, encouraging equal, unbiased treatment of all qubits. \model learns to map physical variables to properties of the ground state of the parameterized Hamiltonian during offline training. Afterwards one can sample from the learned distribution conditioned on previously unseen physical variables to collect measurement records and employ post-processing to predict properties of unknown quantum states. We evaluate \model on large-scale QPE tasks using classically simulated data on the 1D anti-ferromagnetic Heisenberg model with the system size up to 100 qubits. Numerical results demonstrate that \model outperforms baseline models, particularly auto-regressive approaches, under conditions of limited measurement data during training and reduced sample complexity during inference.

Published as a conference paper at ICLR 2025
QUADIM: A CONDITIONAL DIFFUSION MODEL FOR QUANTUM STATE PROPERTY ESTIMATION
Yehui Tang', Mabiao Long', Junchi Yan
Sch. of Computer Science & Sch. of Artificial Intelligence, Shanghai Jiao Tong University
SShangaai Artificial Intelligence Laboratory (yehui trngu,tang, a 1 bong10 yanggg000,annnnmmmaaa junchi jejjtu edu.edu.c.
ABSTRACT
Quantum state property estimation (QPE) is a fundamental challenge in quan- tum many-body problems in physics and chemistry, involving the prediction of characteristics such as correlation and entanglement entropy through statistical analysis of quantum measurement data. Recent advances in deep learning have provided powerful solutions, predominantly using auto-regressive models. These models generally assume an intrinsic ordering among qubits, aiming to approxi- mate the classical probability distribution through sequential training. However, unlike natural language, the entanglement structure of qubits lacks inherent ordering, hurting the motivation of such models. In this paper, we introduce a novel, non-autoregressive generative model called QuaDiM, designed for Quantum state property estimation using Diffusion Models. QuaDiM progressively denoises Gaus- sian noise into the distribution corresponding the quantum state, encouraging equal, unbiased treatment of all qubits. QuaDiM learns to map physical variables properties of the ground state of the parameterized Hamiltonian during offline training. Afterwards one can sample from the learned distribution conditioned on previously unseen physical variables collect measurement records and employ post-processing to predict properties of unknown quantum states. We evaluate QuaDiM on large-scale QPE tasks using classically simulated data on the ID anti-ferromagnetic Heisenberg model with the system size up to 100 qubits. Nu- merical results demonstrate that QuaDiM outperforms baseline models, particularly auto-regressive approaches, under conditions of limited measurement data during training and reduced sample complexity during inference.
INTRODUCTION
Quantum state property estimation (QPE) is a pivotal and challenging problem in the field of quantum many-body physies (Carleo & Troyer, 2017; Torlai et al., 2018; Carrasquilla, 2020; Gebhart et al., 2023; Miles et al., 2023) and chemistry (Kandala et al., 2017: Schütt et al., 2019; Cao et al., 2019; Sajjan al., 2022: Barrett al., 2022). The primary goal is to accurately infer certain properties of the quantum state such as correlation and entanglement entropy using as few quantum measurements as possible. QPE has wide-ranging applications in quantum computing (Dunjko & Briegel, 2018 Torlai & Melko, 2020), quantum cryptography (Brakerski & Shmueli, 2019; Ananth et al., 2022), and quantum simulation (Jia et al., 2019; Schmitt Heyl, 2020). the number of qubits in a system grows, the dimensionality of the corresponding quantum state increases exponentially, making property estimation increasingly difficult. Traditional techniques, like quantum state tomography (D' Ariano et al., 2003; Jullien et al., 2014), become impractical for large systems due to their exponential computational overheads.
Recent advancements in machine learning have opened new avenues for QPE, leveraging the models to extract intricate features from quantum measurement records. An overview about the pipeline is provided in Fig. 1. A notable way relies auto-regressive models, which assumes inherent order among qubits (e.g. left-to-right for 1D chain, or ZigZag order for 2D lattice) and decompose
"Correspondence author. Work was partly supported by NSFC (62222607), Shanghai Municipal Science and Technology Major Project (20215HDDZX0102).

---

Published as a conference paper at ICLR 2025
Reyest Mullape Tanes
Leaming Model
Quantum voolutionn with Lumiltonio comoiiiin (lsy
Samples
Imput Embedding
Past-Proccssing (L.. classical shadow)
Correlation Function
Entanglement Entrupy
Local Observable
Reconstricted Distribution
Physical Variable
Initial
Massurement
State
Distribution
On QPUCCUU
On GPUICPU
Figure 1: An Overview of Quantum State Property Estimation. Given the specified Hamiltonian (x) parameterized by x. the quantum system initiates from state (&aa) and evolves controlled by quantum processing unit (QPU) simulations on CPUs. Measurements are then performed on this final state, generating records stored in classical memory. The machine learning model is trained using measurement records conditioned on and the model's s output would reconstruct the distribution of (e)). The properties of unseen states conditioned on new x such as fidelity can be obtained by post-processing the samples generated by the learned model.
(x) in either physical experiments
their joint probability distribution into conditional probabilities. These approaches include RNN- based models (Carrasquilla et al., 2019; Hibat-Allah et al., 2020; Barrett et al., 2022), CNN-based models (Wu et al., 2019; Sharir et al., 2020) and transformer-ban based models (Cha et al., 2021; Wang et al., 2022; Du et al., 2023; Tang al., 2024a;b).
In auto-regressive models, each qubit conditioned on the states of previously modeled qubits, introducing an inherent bias due to the imposed order. While this method works well in domains where data naturally follows a sequential structure, such as language (Yang et al., 2019; Brown et al., 2020) or time series (Hewamalage et al., 2021), it comes with significant limitation when applied to quantum systems, where qubit interactions are non sequeeeutttaal and the complex entanglement does not adhere to a predefined order. Based on our experimental results, the reliance on sequential modeling could lead to oversimplified representations of the entanglement structure, limiting the model's ability fully capture the intricate correlations between qubits. Consequently, modeling quantum systems with sequential dependencies introduces biases that may fail to capture the true complexity of the entanglement and correlations, leading suboptimal estimations of quantum properties.
To address these challenges, we propose QuaDiM (Quantum state property estimation using Diffusion Models), a novel non-autoregressive generative model based on diffusion models (Ho et al., 2020; Dhariwal & Nichol, 2021; Song et al., 2021a;c). Unlike auto-regressive approaches, QuaDiM treats all qubits equally without imposing a predefined order, thereby avoiding potential biases. The model learns mapping from physical variables, i.e. the parameters of system's 's Hamiltonian, to the ground-state properties associated with those variables. Once trained, QuaDiM capable of generalizing previously unseen quantum systems, allowing it to predict quantum state properties for systems not encountered during training. This generalization capability is particularly valuable in practical scenarios where the quantum states of interest may cannot be prepared on modern noisy intermediate-scale quantum (NISQ) hardware, or quantum states difficult to simulated classically.
The underlying mechanism of QuaDiM is rooted in diffusion models, which iteratively denoise Gaussian noise to generate samples from the target distribution (Austin et al., 2021; Kong et al., 2021). In contrast with auto-regressive models, QuaDiM offers an advantage by treating all qubits equally throughout the denoising process, encouraging to remove the need for any predefined ordering of qubits. Mathematically, QuaDiM initiates from an random Gaussian distribution, which is progressively refined through a series of denoising steps to approximate the true distribution of the quantum state. This iterative process facilitates QuaDiM to capture correlations between qubits more effectively and without the biases introduced by sequential models. By incrementally learning to reverse the noise process, the model naturally handles the complex interactions and entanglements inherent in many-body quantum systems. As a result, QuaDiM not only achieves higher accuracy in QPE but also demonstrates scalability when applied to large-scale systems up to 100 qubits. The model's 's iterative refinement mechanism encourages that all parts of the system considered simultancously. with the hope of being helpful for estimating properties like correlation and entanglement entropy, where the relationships between all qubits are equally important.
2

---

Published as a conference paper at ICLR 2025
We validate the efficacy of QuaDiM through extensive numerical experiments predicting the correlation and entanglement entropy of the ID anti- ferromntroma Heisenberg model with up to 100 qubits. The results show that QuaDiM outperforms baselines, especially auto-regressive approaches, in scenarios with limited training measurements data and reduced sample complexity. QuaDiM'ss ability to handle these challenging conditions highlights its potential for practical applications in quantum computing and related areas, where resource constraints and data scareity are prevalent.
Our contributions. 1) We introduce QuaDiM, the first-ever (to the best of our knowledge) non- autoregressive conditional generative model for QPE. By utilizing diffusion models, QuaDiM itera- tively denoises Gaussian noise to accurately approximate the distribution of unknown quantum states from limited measurements, encouraging equal and unbiased treatment of all qubits. 2) We classically simulate relatively large-scale quantum systems with up to 100 qubits to generate extensive training and test datasets for evaluation, showing QuaDiM's scalability and practical applicability. 3) We conduct experiments on QPE tasks involving the prediction of correlation and entanglement entropy. Our results show that QuaDiM achieves lower prediction errors compared to baselines, particularly outperforming auto-regressive models such as Carrasquilla et al. (2019) and Tang et al. (2024a).
RELATED WORK AND PRELIMINARIES
Learning-based Quantum State Property Estimation. In the domain of learning-based QPE, autoregressive models such as RNN-based (Carrasquilla et al., 2019; Hibat-Allah et al., 2020), CNN- based (Wu et al., 2019; Sharir et al., 2020), and transformer-based models (Cha et al., 2021; Wang et al., 2022; Du et al., 2023; Tang et al.. 2024a) have been extensively explored. These models inherently assume a sequential ordering of qubits, introducing biases that could limit their ability to capture the complex entanglement and correlations in quantum systems where qubits interact non-sequentially. This sequential dependency may lead to oversimplified representations of the entanglement structure, hindering accurate property estimation.
Alternatively, generative models that learn variational wave functions-such as those based on deep Boltzmann machines (DBMs) (Gao & Duan, 2017: Nomura et al., 2017), variational autoen- coders (VAEs) (Rocchetto et al., 2018), and generative adversarial networks (GANs) (Ahmed et al., 2021) have been proposed to represent quantum states. While capable of capturing complex proba- bility distributions, these models typically learn representations for a single specific quantum state and lack the ability to generalize unseen states, restricting their applicability in predicting properties of new quantum systems. Although GANs (Ahmed et al., 2021) could, in principle, generalize to unknown quantum states, their learning objective involves reconstructing the density matrix. This significantly hinders scalability because the complexity of fully describing a density matrix increases exponentially w.r.t quantum system size, making such approaches impractical for large-scale systems.
Quantum State and Quantum Measurement. Quantum computing has the potential to foster advancements in a variety of fields, such as chemistry (Schütt et al., 2019; Cao et al., 2019; Lee et al., 2023), combinatorial optimization (Sanders et al.. 2020; Ye et al., 2025), and machine learning (Bi- amonte et al., 2017; Sajjan al., 2022; Tang & Yan, 2022; Xiong et al., 2024). A quantum bit, or qubit, serves as the fundamental unit in quantum computing. The collection of all qubits within a (sub)system constitutes the quantum state. Prior to measurement, a qubit exists in a superposition of states, but upon measurement, it collapses into definite state. The mathematical representation of quantum state depends on the choice of basis states. For instance, using the orthogonal computational basis states 0)
linear combination (o) a: (0) + 3 1) 3 in the complex vector space C², where a, 3 E C are complex amplitudes satisfying the normalization condition |a|2 + [ 1. An alternative representation of quantum state employs the density operator or density matrix. For example, the density matrix corresponding to |0) is given by Po= 00 (0) = () where (0) denotes the conjugate transpose of [0).
[b] and [1) [ single qubit can be expressed
A general quantum state comprising L qubits can be represented by a wave function:
Voo....., o1)0,,... OL),
(1)
F==1
where V ZL + C assigns a complex amplitude to each fixed configuration = (o1,..., ..h.....000 o1) of the I. qubits, satisfying the normalization condition ΣX= Σx. ((),,...,, 1. Each
3

---

Published as a conference paper at ICLR 2025
o, {l...., }} represents one of the K possible outcomes when a measurement is performed on the i-th qubit. The wave function resides in a complex Hilbert space, with the vector representation (1b) E "C' and its density matrix (v) ((bl E C x both of which grow exponentially large with increasing L. In this paper, we consider the Pauli-6 measurements such that K" = 6.
Quantum measurements translate aspects of quantum information into classical data for subsequent processing, utilizing a set of measurement operators { that satisfy the completeness relation Σ. O₂ = I, where K denotes the total number of possible outcomes. Measuring a qubit yields different results corresponding to the indices of the measurement operators. Specifically, for a quantum state p. the probability of obtaining outcome k upon measurement is p(k) tr(pOx). In systems with L qubits, it is common practice to perform measurements on all qubits simultaneously or in parallel (Leibfried et al., 1996; Jullien et al., 2014). According to Born's ''s' rule (Born, 1926) in quantum mechanics, this measurement process produces a string of outcomes o = (O).
where each o; E (I,..., K}. with probability v (oi....,, TT)) defined in Eq. 1.
Diffusion Models. Diffusion models are a class of generative models that have gained significant attention in recent years due to their ability to produce high-quality samples from complex distribu- tions (Ho et al., 2020; Song et al., 2021a). These models typically consist of two main processes: a forward process that gradually adds noise to the data, and a reverse process that aims to reconstruct the original data from the noisy representation. In the forward process, a data point sampled from a real-world distribution, denoted as Zu ~ p(z). is progressively corrupted by adding Gaussian noise. This process transforms Zn into standard Gaussian noise vector Zr .N. I) over a predefined number of steps T. For each step E (1, ...,, T}. the perturbation is governed by the conditional distribution /(,,/ 1) ((,;; v1 1, 3, where Be E (0, 1) represents varying variance scales at each time step. This formulation allows for a controlled degradation of the original data, enabling the model to learn to navigate through the noise space effectivelyy Once the forward process is complete, the reverse denoising process seeks to recover the original data Zo by sampling from Zy. This is achieved through the training of diffusion model fe. which learns the conditional distributions necessary to reverse the noise addition process. The learning objective typically involves minimizing the difference between the predicted and actual distributions, often through techniques such as score matching (Song et al., 202 or variational inference (Sohl-Dickstein et al., 50. Diffusion models have demonstrated remarkable capabilities in generating across various domains, including image synthesis (Rombach et al., 2022), audio generation (Kong et al., 2021), and more recently, applications in solving combinatorial optimization problems (Li et al., 2023) and imitation learning (Pearce et al., 2023). Their inherent flexibility and robustness also make them an ideal candidate for capturing the complex distributions associated with quantum systems, thus paving the way for advancements in quantum machine learning applications.
samples
Classical Shadow for Post-Processing A standalone generative model and the samples drawn from its distribution cannot directly reconstruct quantum state or extract certain features of it. Post-processing is necessary. In this paper, we consider using classical shadow (Huang et al., 2020) based on randomized single-qubit measurements to predict quantum state properties. This measurement procedure is highly efficient on both quantum experiments and classical simulations. Accordingly, we employ Pauli-6 positive operator valued measure (POVM) to collect discrete mea- surement records, which are used to train the generative model. Concretely, the Pauli-6 POVM is given as O = { x (0),00, x [1) X [+) x |-X + x (r), $ x [IX(I], where {0),, [1)), (1)), ). {Ir),
{2)} stand for the eigenbasis of the Pauli operators X, and Y. re- spectively. Suppose that the quantum system has I. qubits, measuring a single qubit using Opauti where i E {1...., L} denotes the i-th qubit and m E (1.....M denotes the m-th measurement. Then the quantum state properties such as corre- lation and the entanglement entropy can be estimated using the LM snapshots. The details can be found in Sec. 4. After the learning model have been trained, we sample from the model's distribution conditioned the unseen parameters and collect the samples, i.e., the measurement data. Then classical shadow is used for post-po post-processing to estimate the propertics of unknown quantum states and the results are compared the ground truth to evaluate the model's performance.
leads to snapshot (!!))) 3|s; (m)) X=!m!
METHODOLOGY
We first define the task of QPE in Sec. 3.1. Key methods and insights underlying QuaDiM are then detailed in Sec. 3.2, where we discuss its mathematical foundations and core innovations.
4

---

Published as a conference paper at ICLR 2025
Forward Process
Reverse Process
Ha( it
p (m
Ps (mu)a)
Lincr Projection
Embedding
Select
LL
Layer Normalization
Records
Multi-head Attention
(0(2;-11 (2))
Rounding
Collect
Noised
Timestop
Lateos
Condition
Ky IO,
Zo
a) Forward and Revere Process of QuaDiM
b) Structure of Transformer
Figure 2: a) The diffusion forward process iteratively perturbs the input by adding Gaussian noise, while the reverse process incrementally removes the noise to recover the original distribution. To facilitate the transition between the latent variable Zo and the observed records σr, embedding function (Li et al., 2022) is employed, along with a rounding step (Gong et al., 2023). b) The main learnable part is multi-layer transformer. It models the mean of the posterior distribution qo-
3.1 PROBLEM DEFINITION
The problem is to learn to predict the quantum state properties ground states parameterized Hamiltonians. We consider family of Hamiltonians, where the Hamiltonian H (x) is parameterized
by a set real parameters X IRP These parameters are physical condition variables determining the evolution dynamics of the quantum system. Different physical conditions could lead different ground states. The neural network is trained on training data consisting of sampled values of {x, {x.] each accompanied by the corresponding measurement records R. ZM,. x obtained from Pauli-6 measurements repeatedly operated M,n times on the ground state p(x,) of (x.). This training data could be obtained from either classical simulations or quantum experiments. During the inference phase, we sample Mout times from the trained neural network conditioned on a new values of x" E {x;; N in the test data and outputs a collection of measurement data R E Mout x L of the unknown quantum state p(x*). Ground state properties can then be estimated using post-proccss post-processing on the Rs such as the classical shadow.
3.2 QUADIM
We propose QuaDiM extend vanilla diffusion models to learn conditional quantum state represen- tations (as shown in Fig. 2), concerning the architecture and the training objective. For conciseness if otherwise specified, we will omit the subscripts of the variables in the subsequent text.
3.2.1 FORWARD NOISING PROCESS WITH QUADIM
Although diffusion models have achieved success in image generation, applying continuous diffusion models to inherently discrete quantum measurement records remains non-trivial. We follow the Diffusion-LM (Li et al., 2022) in the text generation domain to design embedding function that projects discrete measurement records into a high-order continuous feature space. Specifically, given physical condition and the corresponding measurement records R. QuaDiM first randomly samples a measurement string a = (ø1,..., a which is a single row of R. To apply a continuous diffusion model discrete measurements, a token embedding function EMB (01) is introduced to map the single-qubit measurement outcome σ to a d-dimensional vector for each / E (1...., L}. The token embeddings are jointly trained with diffusion model's parameters. We empirically find that these learnable embeddings can achieve better numerical results than the fixed embeddings. The token embedding is accompanied with trainable positional embedding (PE) with the same hidden dimension d. Thus the token embedding function is written as EMB(o) = (EMB(01) PE(1). (F(1),..... EMB( PR([.) E L The initial input to the diffusion model in time step is an additional Markov transition from EMB(O), given
Po (2olo) = V (zo; EMB(O),
2
where 8 is the set of parameters of learnable token embeddings and positional embeddings. The forward noising process gradually corrupts zo into a standard Gaussian noise ZT ~ N (0, I). For each time step E (1....,T), the perturbation is defined as p(z (z) = N (z:: B.
5

---

Published as a conference paper at ICLR 2025
3.2.2 REVERSE CONDITIONAL DENOISING PROCESS WITH QUADIM
The proposed conditional denoising employs classifier-free approach, i.e. we do not need to train an additional classifier to navigate the denoising process. Instead of performing conditional denoising directly on the discrete quantum measurement records, we perform it on the sequence of continuous latent variables Zo.T defined by the QuaDiM. It enables a simple gradient-based algorithm to perform complex, controllable quantum state generation based on diffusion models. Specifically, conditional
denoising is equivalent to decoding from the posterior g0 (Zo:r x) = g(²r) II:- yo (2--1 22, x). For each denoising step. we model as a Gaussian qo [Z(-1/Z(. = N (Z Me Z1, x, t), 7 (t)l), where Hi and (t) are the predicted mean and variance:
B
3e 2t
BIV
3 So x,,, 7(1)
3,3
1 b2, x, 1)
B.
B.
3.
where 3, = ITi=-( 3c).
The network po(2,, x, 1) is akin to the multi-layer transformer decoder (Vaswani et 2017). The distinction is that an additional time step embedding is specified by the transformer sinusoidal position embedding (PE) (Ho et al., 2020). Besides, to achieve the controllable generation of quantum states conditioned on the physical variable x, we use feed-forward network (FFN) with one hidden layer to transform x into a hidden feature with the same dimension d. This feature is viewed as a global information added to the token embeddings of measurement records, such that the imput to the transformer at time step t is H 2, + +FFN)))FF(x)) Suppose that the (I 1)-th layer's 's output H E Lxd where is the hidden dimension, we have:
H' MULTTHEADATTENTION
[0
0% W
with 0 (H'-'wa 4. )
W
)
T
H W
vd
where [.] denotes the concatenation operation, W! q.j* W W E Rdxd/h and W Rdxd are the parameters to be learned. In this paper. all the experimental results of QuaDiM are reported for a transformer configuration consisting of heads, layers, and 128 hidden dimensions. The maximum denoising time steps is set to T = 2000.
3.2.3 LEARNING OBJECTIVE
The canonical learning objective is to minimize the KL divergence between the joint probability distribution p( pZoor (r) and go (/) corresponding to the forward and reverse processes defined by the diffusion model, respectively. The objective given as
p p(Zo.T ④))
CKL = (p(xor)(or) log
(5)
go (x:r(x)
where p(zo:r| (o) 58 (2o) or) IT p(z, 1 denotes the distribution of the forward process, and go (x) = ((T)) [[= go (2 Z x) represents the distribution of the reverse process. Note that in the forward process, Ps (z))or) is the Markov transition (as given in Eq. 2) from the discrete quantum measurements to their continuous embeddings, allowing the token embeddings and model parameters jointly participate in gradient descent. The objective can be further simplified
LKL = (ot log p(ZT(20)
T
p(z.- 20. 2x) + log (2 (z))
Lo (Zo ④r)
+ log
Σ 2-2
q(Zr)
g0 (0(21-1 lz,)
The detailed derivation and explanation of each term are given Appendix A. In the reverse process, the diffusion model iteratively removes Gaussian noise Zr and outputs the embeddings Zg. To estimate the reconstruction quality from the model's output distribution when training the model, a rounding step (Li et al., 2022; Gong et al., 2023) is added to map the continuous obtained from the denoising process back discrete measurement outcomes a of length L.. Specifically. the step is given as jo (o(aa) = I Āo(a), (2o.) and Zo,1 is the I-th component of Zu E IR. Lxd Now we introduce the overall objective the combination of two terms: maximizing the average negative log-likelihood Epizaixxe ..) log ğo (a (o/z)) and minimizing the KL divergence LK Thus the objective as follows, where the detailed derivation is given in Appendix B.
min L = min LKL - Eppno| xx)))) log qo (a [20) 6.0
8,0
T
(7)
min (x, 6.0
[|EMB(a)- (z +
||z0 vo(z,,x,t))) - log
GER

---

Published as a conference paper at ICLR 2025
Table 1: RMSE of predicting the correlations of all subsystems of size two on the test dataset. The
result is averaged over Heisenberg model instances and each pair of adjacent qubits. For CS, M denotes the number of input measurements. While for the neural network-based approaches M denotes the number of sampled measurements Mout from trained models. The best results are emphasized in red while the second-best best results are distinguished in blue.
L-
10
L-
30
1 70
Z- 100)
Mathol
M/ - 100
1000
10000
20000
100
1000
10000
20000
100
1000
10000
2000
100
1000
10000
20000
0,1564
0,0599
0,0156
0,0077
0.1696
0,5388
9.1133
0.0121
0.1771
0,0545
0,1722
0,1211
0.1724
0.547
0,0122
0.122
RBEK NTK
0.796
0,0639
0.5788
0.0491
00775
0,0622
0.8655
0,0470
RNN LLMAQPE
0,1328
0/152
0.0155
011119
0.3795
0.0671
010164
0.01IN
0.2137
0,0399
0.0240
0,1133
U.2325
0U0006
0.0251
00165
0,1316
0,0489
0.0136
0,00M3
0.1624
1J8513
0J1142
0.0097
0.1814
0.0527
0.0155
0.8116
0L1759
0J1531
0.0152
0.0144
Ouss
0.1369
0,0432
0,0097
0,0085
0.15K2
IJ9465
IJI133
0,0091
0.1679
0,0473
0,8117
0.8092
IL1686
0J1478
0.0125
0.0998
4 EXPERIMENTS
In this section, we first introduce dataset construction and baselines in Sec. 4.1 and Sec. 4.2. In Sec. 4.3 and Sec. 4.4, we evaluate QuaDiM's performance on QPE tasks compared to state-of-the-art baselines, focusing predicting correlation and entanglement entropy in various system sizes.
4.1 DATASETS
We focus on investigating the ground state the one-dimensional antrmeerroaannctic Heisenberg model, as it can be efficiently simulated classically relatively large scale. Its Hamiltonian is parameterized by x E RL-1 and can be written as
H(x) = Σ xy (X,X++1 + Y, + Zz
(
where L is the number of qubits and x is sequence of real numbers for the coupling strength among qubits. X, Y, Z denotes the Pauli operator and X, indicates that Pauli X is operated on the i-th qubit.
For the sake of simplicity, we first describe how single sample is generated, with other samples in the dataset being produced in a similar manner. Given the length T. of the quantum system, we first randomly sample I, - distinct values of x to construct the corresponding Hamiltonian H(x). We then obtain the ground state of the H (x) using classical simulation based on tensor networks (Fishman et al., 2022). Thanks to the efficiency of tensor networks, we are able to simulate quantum systems from 10 to 100 qubits on a CPU cluster. Subsequently, Pauli-6 positive operator-valued measure (POVM) is utilized to measure the ground state corresponding to H (xx) and the output is string- like sequence o = (aI,...,,o TT.) where o E (1, K}. The same measurement procedure is performed Man. times independently on the Min. copies of the ground state. As a result, the model input consists of x along with the corresponding measurement records R E 2 xL, The above process is repeated N and Noe times to obtain the training dataset (x; N with measurement records (R. NT and test dataset (x, with (R,)
As for the ground truth labels, which represent the ground state properties of the Hamiltonian (xx such as correlation and entanglement entropy, we utilize exact diagonalization (WeiBe & Fehske, 2008) for quantum systems with L 10 to obtain the true labels. For quantum systems with L L> 10, we collect a substantial amount of measurement data by setting M = 320 000 and approximate the corresponding labels using classical shadow (Huang et al., 2020).
4.2 BASELINES
We consider the classical shadow (CS) (Huang et al., 2020) a SOTA learning free protocol for constructing the representation of an unknown quantum state. For learning-based baselines, we consider kernel methods including the Radial Basis Function Kernel (RBFK) and Neural Tangent Kernel (NTK) implemented Huang et al. (2022). We also explore advanced deep learning approaches, including Recurrent Neural Network (RNN)-based model (Carrasquilla et al., 2019), and a transformer-basee SOTA model LLM4QPE (Tang et al., 2024a).
For all the methods, we set 'rr 100 and Ne = 20, with the number of qubits in the quantum system L E (10,40 40, 70, 1003. To construct the training set, we perform repeated measurements of M,. 1000 for each ground state. To evaluate the performance of the trained QuaDiM and auto-rees regressvgee baselines in predicting certain quantum state property w.r.t. different sampling counts, the model is sampled M times, and the outputs are processed to predict the ground state property
7

---

Published as a conference paper at ICLR 2025
RONESE-0,32300
RMSE-0221
RMSEE-01316
RMSEE RMSB--00133
L.MEQPE
Random Hriscaheng model sct
535
QuDM
-38
Mace 100
Maat
Maxe 10000
Figure Visualization predicted correlations Cij for the ground state of the anti-ferromagnetic Heisenberg model of length L 10 with different number of samples Mout from the trained models. The chain in the upper right corner represents a Heisenberg model selected from the test set, where the width of the edges indicates the coupling strength F;. The upper part shows the prediction results of LLM4QPE (Tang et al., 2024a), while the lower part presents those of proposed QuaDiM alongside the ground truth. For all the settings, we fix Min 1000 and use classical shadow for post-processing.
RMSB--040
RMSE--1177
RMSE-L0887
RMSE-400114
RMSE-00199
5. QaIM
[ Gausson Noiee
400
C 800
t- 1200
t 1600
2000, autdlas)
Figure 4: Visualization of predicted correlations of QuaDiM Č₁2 for the ground state of the ID 10. The measurement samples are generated
antomaerroaagnttic Heisenberg model of length from the different denoising time step f. For all the settings, we fix Min = 1000 and Mout = 10000.
of the Hamiltonian in the test dataset. We use the Root Mean Square Error (RMSE) to assess the difference between the predicted property and the true one, with smaller RMSE indicating that model achieves better predictive accuracy with fewer measurement samples or sampling iterations. To adapt our proposed generative model, we made appropriate adjustments to all auto-regressive baselines including RNN and LLM4QPE, to facilitate generative training and sampling from the trained model, while employing the same classical shadow protocol for post-processing to ensure a fair comparison. For the kernel methods, since generative training is not possible, the model input retains only the physical condition x. For specific details about baselines, please refer to Appendix D.
4.3 PREDICT THE CORRELATION
The first learning task is predict the correlations of all subsystems of size Denote the density matrix of the quantum state as p. For two qubits located at different positions i and j. the correlation described as C.j- tr(O;;p) where O.ş (X,X, + Y,Y, Z,Zz). Given physical condition x, the model along with post-processing is learned to approximate the mapping f
Cij-
The data sampled from the trained model will be fed into the classical shadow for post-processing. Specifically, assuming the number of samples is Mout- we obtain an output Rout E Z Man x where
each element can be converted into snapshot p (m))
(m) (s(m)
where {l...., L} and
m E {1...., Moue), J. as discussed in Sec. 2. Then the predicted correlation between the i-th qubit and the j-th qubit is calculated by
Mout
Cry =
1
Σ
tr
(p! (m)
pfmm)
(o: 8 803)
(9)
3.Mout (E(x.N.2) X
As shown in Tab. 1, QuaDiM achieves the lowest RMSE across all the configurations, showcasing its ability to deliver precise correlation predictions even with a moderate Mout- An important advantage of QuaDiM is its ability to maintain excellent performance even with fewer Mout. For instance, at Mous = 100, QuaDiM's RMSE of surpasses CS and remains competitive with the auto-regressive baselines. This trend persists acToss larger system sizes as well. QuaDiM demonstrates strong performance at low sample counts, indicating its efficiency in utilizing limited samples from the trained model for accurate quantum state property estimation.
8

---

Published as a conference paper at ICLR 2025
As expected, the RMSE decreases with increasing Mout, reflecting that all models except for kernel methods improve in predictive accuracy as more measurement samples become available. However, the degree of improvement varies significantly between methods. The CS protocol performs reasonably well for small Mout but shows diminishing returns at larger Mout, particularly for larger system sizes (T. = 40 and beyond). In contrast, QuaDiM continues to benefit from increased Mout. which suggests that the model is better at learning and generalizing from larger quantum systems.
Kernel-based methods (RBFK and NTK) show competitive performance for medium system sizes (L = 40), but they struggle to match the performance of the generative models (QuaDiM, LLM4QPE and RNN) at higher sample sizes or larger system sizes. This could be due to the inherent limitations of kernel methods in capturing the complex structure of large quantum systems, which generative models like QuaDiM can better model through learned representations.
One of the most critical challenges in quantum state estimation is scalability. While some methods, such as the classical shadow approach, demonstrate robust performance for smaller system sizes (L = 10), their performance becomes inferior compared to the neu-
L 100
0.250
0.250
RMSE
0.062
0,062
0016
0.016
100
250 s00 1000
50
100
250
1000
ral network- brrkk- methods as L increases. This is consistent with the findings in Huang et al. (2022) that neural networks could more easily display non- capture
Truining size of Min
Training size Min
RNN
LLM4QPE
QuaDiM
local correlations, allowing in principle quantum states with higher entanglement for rela- tively larger quantum systems. Notably, for L = 100, the RNN-based baseline yields significantly higher
Figure 5: RMSE of predicting the correla- tions of all subsystems of size two the test dataset of the regrereegeeevve baselines and the proposed QuaDiM for L (70, 100). The models are trained using different of num- ber of input measurements Min.
RMSE values compared to others, indicating that the auto-regressive training and generation mechanism makes oversimplified representations of the entangle- ment structure, limiting the model's ability to fully capture the intricate correlations between qubits. Conversely, QuaDiM shows resilience to scaling, maintaining low RMSE values across all system sizes, with only slight increase in RMSE L increases. This scalability suggests that QuaDiM is particularly well-suited for handling large quantum systems, making it promising approach for real-world quantum computations where the number of qubits can be extensive (for example, IBM's latest commercial quantum computer has more than 1,000 qubits).
More results about the performance w.r.t. the samples complexity of the trained models. The experimental results visualized in Fig. show the predicted correlations Cij for the ground state of the ID anti-ferromagnetic Heisenberg model, with varying sample sizes Mout. QuaDiM outperforms the auto-regressive baseline LLM44QE in achieving lower RMSE with fewer samples. As the sample size increases, QuaDiM continues to demonstrate superior performance, achieving an RMSE of 0.0097 at Mout 10000, compared to LLM4QPE's 0.0136. We also illustrate QuaDiM's 's predicted correlations at various denoising time steps in Fig. 4. As the denoising process progresses, the RMSE steadily decreases. This demonstrates that QuaDiM effectively reconstructs the distribution of quantum states as the denoising process advances, with each time step reducing the noise and improving the fidelity of the sampled quantum states. A more tineggraindd visualization of the prediction performance of QuaDiM is depicted in Fig. 6 in Appendix E.
More results about the performance w.r.t. the number of measurements used for training. The experimental results shown in Fig. 5 highlight the predictive error of auto-regressive baselines and our proposed QuaDiM when predicting correlations for subsystems of size two within quantum systems of T. {70, 100). The models are trained using different numbers of input measurements Min ranging from 50 to 1000. QuaDiM consistently outperforms the auto-regressive baselines across all training sizes. For both L E {70, 100), achieves lower RMSE with fewer input measurements compared to the other methods, showing a advantage in training efficiency. Notably, as Mnn increases, QuaDiM RMSE decreases more sharply than that of the baselines, indicating that can better approximate the (classical) probability distribution of quantum states using fewer input measurements. quantum computation, where measurement resources are often limited for both quantum experiments and classical simulation (Gebhart et al., 2023), QuaDiM's 's ability to perform well with fewer measurements offers a significant advantage over rugo-regressive methods.

---

Published as a conference paper at ICLR 2025
Table 2: RMSE of predicting the entanglement entropies of all subsystems of size the test dataset. The result is averaged over Heisenberg model instances and each pair of adjacent qubits. The best results are emphasized in red while the second-best results are distinguished in blue
1.-10
7-- 10
1.
70
7- 100
Medhod
11 - 100
1000
10000
20000
100
1000
10000
2000
100
1000
10000
20000
100
1000
10000
20000
0,9066
0.0922
0,0004
0,0119
0.487
0.9977
0.0244
0,0299
0.6421
0.0943
0.0312
0,0288
0.6518
0.0998
0.0357
0/0316
RBFK NTK
0.1268
0.1057
0.9997
0,0752
0.1379
0.1034
009953
0,0719
RNN
0,5255
0.1164
0,0187
0,0115
0.6322
0.1054
0.2266
0.0212
0,7948
0.1305
0.0814
0,0855
08229
0.1476
0.0517
0.049
LLM-OPE
0,4937
0,1948
0.0176
0,0102
0.5M78
1.0996
0JI233
0.0307
0,6258
0,0912
0.8276
0/8251
IL6392
0.1055
0,0312
0.0266
Ours
0,5479
0,0867
0,0132
0,6089
0,5629
0.0661
0.0187
6.0145
0.8970
0.0879
0,0245
0,8218
0.6128
0.0928
0.0281
0.0243
4.4 PREDICT THE ENTANGLEMENT ENTROPY
Another learning task is to predict the second- order Rényi entanglement entropy of the subsystem A, which is formulated A system. The required number of measurements scales exponentially in the size of the subsystem A, but is independent of total system size (Huang al., 2020). In this paper, entanglement entropy of all the subsystems of size two. This leads to a local quadratic feature for a unknown quantum state p. which can be also efficiently estimated by the classical shadow given by
log(tr (p²) )), where A is subsystem of the n-qubit quantum
consider learning the
Eig = - -ogg(t((PAA))) - log(tt(SSA,)))))) A., 2 p)) = log(tr(S A. BlA @ BlÅzl))
Σ Ar(SA pmm) & ( = E,j,
(10)
Mout (Mous 1j
where #1. و are two independent snapshots of the unknown state p and SA, denotes the local swap operator of two copies of the subsystem comprising the i-th and the j-th qubits. The [[p] represents the averaged value among Mous samples from the trained model. Given a specific physical condition x, the model along with post-processing is learned to approximate the mapping f x Esj.
Tab. 2 compares models' performance in predicting the entanglement entropy of subsystems of size two. The key metric is also the RMSE, with lower values indicating better predictive accuracy. The results are presented across different system size I. and sample sizes Mout from the trained models. CS method performs reasonably well lower system sizes and sample counts, but it struggles to maintain competitive RMSE values as both M and L increase. Similarly, the kernel-based methods demonstrate some degree of competitiveness but fall behind neural network- based models at larger system sizes and higher sample counts. It can be seen that QuaDiM achieves superior performance across most system sizes and sample counts, consistently achieving lower RMSE values compared to baseline models. This trend continues as the sample size increases, with QuaDiM maintaining its advantage over the other models. The RMSE reduction demonstrates that QuaDiM can accurately predict the entanglement entropy even with relatively small number of samples.
Performance across system sizes. QuaDiM also scales effectively as the system size increases. This indicates that QuaDiM could be practical choice for real-world quantum computations where the number of qubits can be large. The model's ability to maintain low RMSE values across different system sizes shows its scalability even for the relative large quantum system with 100 qubits.
Sample efficiency from the trained models. One of the most important advantages of QuaDiM its sample efficiency. It consistently achieves lower RMSE values with fewer samples for post-processing As the number of samples increases, the performance gap widens, with QuaDiM achieving an RMSE of 0.0145 with Mout 20, 000, while LLM4QPE and CS exhibit higher RMSE. This sample efficiency is critical in post-processing, as large samples leads to prohibitive post-processing costs.
5 CONCLUSION
In this paper, we introduced QuaDiM, novel non-autoregressive model for QPE based diffusion models. Our results demonstrate it achieves superior performance in predicting quantum properties such correlation and entanglement entropy with limited measurement data and reduced samples.
Limitations and Future Work. While QuaDiM demonstrates promising improvements, there is still room for further investigation regarding its scalability to highly entangled quantum systems and its sampling efficiency. Future work will aim to extend the model to accommodate more complex quantum systems and incorporate advanced methods to enhance the speed of diffusion sampling.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Shahnawaz. Ahmed, Carlos Sánchez Muñoz, Franco Nori, and Anton Frisk Kockum. Quantum state tomography with conditional generative adversarial networks. Physical review letters, 127(14): 140502, 2021.
Prabhanjan Ananth, Luowen Qian, and Henry Yuen. Cryptography from pseudorandom quantum states. In Annual International Cryptology Conference, pp. 206-236. Springer, 2022.
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusson models in discrete state-spaces. Advances in Neural Information Processing Systems, 34: 19981-179999, 2021.
Thomas D Barrett, Aleksei Malyshev, and AI Lvovsky. Autoregressive neural-network wavefunctions for ab initio quantum chemistry. Nature Machine Intelligence, 4(4):351-358, 2022.
Jacob Biamonte, Peter Wittek, Nicola Pancotti, Patrick Rebentrost, Nathan Wiebe, and Seth Lloyd. Quantum machine learning. Nature, 549(7671):195-222, 2017.
Max Born. Quantenmechanik der stoBvorgänge. Zeitschrift für physik, (8(11-1)):003---81811 1926.
Zvika Brakerski and Omri Shmueli. (pseudo) random quantum states with binary phase. In Theory of Cryptography Conference, pp. 229-250. Springer, 2019.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 13:187771991, 2020.
Yudong Cao, Jonathan Romero, Jonathan P Olson, Matthias Degroote, Peter Johnson, Mária Kieferová, Ian D Kivlichan, Tim Menke, Borja Peropadre, Nicolas PD Sawaya, et al. Quantum chemistry in the age of quantum computing. Chemical reviews, 119(19):10856-10015,, 2019.
Giuseppe Carleo and Matthias Troyer. Solving the quantum many-body problem with artificial neural networks. Science, 355(6325):602-606 2017.
Juan Carrasquilla. Machine learning for quantum matter. Advances in Physics: X, 5(1)11997528
2020.
Juan Carrasquilla, Giacomo Torlai, Roger G Melko, and Leandro Aolita. Reconstructing quantum states with generative models. Nature Machine Intelligence, 133:155--111, 2019.
Peter Cha, Paul Ginsparg, Felix Wu, Juan Carrasquilaa, Peter L McMahon. and Eun-Ah Kim. Attention-based quantum tomography. Machine Learning: Science and Technology, 3(1)001LT01, 2021.
Mauro D' Ariano, Matteo GA Paris, and Massimiliano F Sacchi. Quantum tomography, Advances in imaging and electron physics, 128:206-309, 2003.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 1118880-8744, 2021.
Yuxuan Du, Yibo Yang, Tongliang Liu, Zhouchen Lin, Bernard Ghanem, and Dacheng Tao. Shad- ownet for data-centric quantum system learning. arXiv preprint arXiv:2308.1 12900 2023.
Vedran Dunjko and Hans Briegel. Machine learning & artificial intelligence in the quantum domain: a review of recent progress. Reports on Progress in Physics, 81(7):07400,, 2018.
Matthew Fishman, Steven R. White, and E. Miles Stoudenmire. The Tensor Software Library for Ten- sor Network Calculations. SciPost Phys. Codebases, pp. 4, 2022. doi: 00211468ssiiotttPyyyCoohhh
4.
Xun Gao and Lu-Ming Duan. Efficient representation of quantum many-body states with deep neural networks. Nature communications, 8(1 8(1):662, 2017.

---

Published as a conference paper at ICLR 2025
Valentin Gebhart, Raffaele Santagati, Antonio Andrea Gentile, Erik M Gauger, David Craig, Natalia Ares, Leonardo Banchi, Florian Marquardt, Luca Pezzè, and Cristian Bonato, Learning quantum systems. Nature Reviews Physics, (3):141-156, 2023.
Shansan Gong, Mukai Li, Jiangtao Feng, Zhiyong Wu, and Lingpeng Kong. DiffuSeq: Sequence to sequence text generation with diffusion models. In International Conference Learning Representations, 2023.
Hansika Hewamalage, Christoph Bergmeir, and Kasun Bandara. Recurrent neural networks for time series forecasting: Current status and future directions. International Journal of Forecasting, 37 (1):388-421) 2021.
Mohamed Hibat-Allah, Martin Ganahl, Lauren E Hayward, Roger G Melko, and Juan Carrasquilla. Recurrent neural network wave functions. Physical Review Research, 2(2):023358, 2020.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-663 6851, 2020.
Hsin- Yuan Huang, Richard Kueng, and John Preskill. Predicting many properties of a quantum 100--1070--15 2020.
system from very few measurements. Nature Physics,
Hsin- Yuan Huang, Richard Kueng, Giacomo Torlai, Victor V Albert, and John Preskill. Provably efficient machine learning for quantum many-body problems. Science, 377 2022.
Zhih-Ahn Jia, Biao Yi, Rui Zhai, Yu-Chun Wu, Guang-Can Guo, and Guo-Ping Guo. Quantum neural network states: A brief review of methods and applications. Advanced Quantum Technologies, (7-8):1800077, 2019.
Jullien, P Roulleau, Roche, A Cavanna, Jin, and DC Glattli. Quantum tomography of an electron. Nature, 51475224460333-67 2014.
Abhinav Kandala, Antonio Mezzacapo, Kristan Temme, Maika Takita, Markus Brink, Jerry Chow, and Jay M Gambetta. Hardware-effficient variational quantum eigensolver for small molecules and quantum magnets. nature, ())2)262)) 2017.
Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. In International Conference Learning Representations, 2021.
Seunghoon Lee, Joonho Lee, Huanchen Zhai, Yu Tong, Alexander M Dalzell, Ashutosh Kumar, Phillip Helms, Johnnie Gray, Zhi-Hao Cui, Wenyuan Liu, et al. Evaluating the evidence for exponential quantum advantage in ground-state quantum chemistry. Nature communications, 14 (1):1952, 2023.
Dietrich Leibfried, DM Meekhof, BE King, CH Monroe, Wayne M Itano, and David J Wineland. Experimental determination of the motional quantum state of trapped atom. Physical review letters, 77(21):4281 1996.
Xiang Li, John Thickstun, Ishaan Gulrajani, Percy S Liang, and Tatsunori B Hashimoto. Diffusion-Im improves controllable text generation. Advances in Neural Information Processing Systems, 35: 4328 4343, 2022.
Yang Li, Jinpei Guo, Runzhong Wang, and Junchi Yan. T2t: From distribution learning in training to gradient search in testing for combinatorial optimization. Advances in Neural Information Processing Systems, 36:50020-50040, 2023.
Cole Miles, Rhine Samajdar, Sepehr Ebadi, Tout T Wang, Hannes Pichler, Subir Sachdev, Mikhail D Lukin, Markus Greiner, Kilian Q Weinberger. and Eun-Ah Kim. Machine learning discovery of new phases in programmable quantum simulator snapshots. Physical Review Research, 5(1): 013026, 2023.
Michael A Nielsen and Isaac L Chuang. Quantum computation and quantum information. Cambridge university press, 2010.
12

---

Published as a conference paper at ICLR 2025
Yusuke Nomura, Andrew S Darmawan, Youhei Yamaji. and Masatoshi Imada. Restricted boltzmann machine learning for solving strongly correlated quantum systems. Physical Review B, 96(20): 205152, 2017.
Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, and Sam Devlin. Imitating human behaviour with diffusion models. In The Eleventh International Conference on Learning Representations, 2023.
Andrea Rocchetto, Edward Grant, Sergii Strelchuk, Giuseppe Carleo, and Simone Severini. Learning hard quantum distributions with variational autoencoders. npj Quantum Information, 4(1):28, 2018.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Manas Sajjan, Junxu Li, Raja Selvarajan, Shree Hari Sureshbabu, Sumit Suresh Kale, Rishabh Gupta, Vinit Singh, and Sabre Kais, Quantum machine learning for chemistry and physics. Chemical Society Reviews, 51(15))6475-657 2022.
Yuval R Sanders, Dominic W Berry, Pedro CS Costa, Louis W Tessler, Nathan Wiebe, Craig Gidney, Hartmut Neven, and Ryan Babbush. Compilation of fault-tolerant quantum heuristics for combinatorial optimization. PRX quantum, (22:022012,,003) 2020.
Markus Schmitt and Markus Heyl. Quantum many-body dynamics in two dimensions with artificial neural networks. Physical Review Letters, 125(10): 2020.
Kristof T Schütt, Michael Gastegger, Alexandre Tkatchenko, K-R Müller, and Reinhard J Maurer. Unifying machine learning and quantum chemistry with deep neural network for molecular wavefunctions. Nature communications, 10(1):5024, 2019.
Or Sharir, Yoav Levine, Noam Wies, Giuseppe Carleo, and Amnon Shashua. Deep autoregressive models for the efficient variational simulation of many-body quantum systems. Physical review letters, 124(2):020503, 2020.
Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representations. arXiv preprint arXiv: 0001122 2018.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference machine learning, pp. 2256-2265. PMLR, 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In Interna- fional Conference on Learning Representations, 2021a.
Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Advances in neural information processing systems, 34:1415-1428, 2021b.
Yang Song, Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021c.
Yehui Tang and Junchi Yan. Graphqntk: Quantum neural tangent kernel for graph data. Advances in Neural Information Processing Systems, 35:6104-6118. 2022.
Yehui Tang, Hao Xiong, Nianzu Yang. Tailong Xiao, and Junchi Yan. Towards l1m4qpe: Unsuper- vised pretraining of quantum property estimation and a benchmar. In The Twelfth International Conference Learning Representation, 2024a.
Yehui Tang, Nianzu Yang. Mabiao Long, and Junchi Yan. Ssl4q: supe-supe learning of quantum data with application to quantum state classification. In Forty-first International Conference on Machine Learning, 2024b.
13

---

Published as a conference paper at ICLR 2025
Giacomo Torkai and Roger G Melko. Machine-learning quantum states in the nisq era. Annual Review of Condensed Matter Physics, 11:325-344 2020.
Giacomo Torlai, Guglielmo Mazzola, Juan Carrasquilla, Matthias Troyer, Roger Melko, and Giuseppe Carleo. Neural-network quantum state tomography. Nature Physics, 14(5):47-450,, 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Haoxiang Wang, Maurice Weber, Josh Izaac, and Cedric Yen- Yu Lin. Predicting properties of quantum systems with conditional generative models. arXiv preprint arX1v:2111.1944,, 2022.
Alexander WeiBe and Holger Fehske. Exact diagonalization techniques. Computational many-particle physics, pp. 529-544, 2008.
Dian Wu, Lei Wang, and Pan Zhang. Solving statistical mechanics using variational autoregressive networks. Physical review letters, 122(8)000002, 2019.
Tailong Xiao, Jingzheng Huang, Hongjing Li, Jianping Fan, and Guihua Zeng. Intelligent certification for quantum simulators via machine learning. npj Quantum Information, 8(1): 8(1):138, 2022.
Hao Xiong, Yehui Tang, Xinyu Ye, and Junchi Yan. Circuit design and efficient simulation of quantum inner product and empirical studies of its effect on near-term hybrid quantum-clasic machine learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 26162-26177 270 2024.
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoe V Le. Xlnet: Generalized autoregressive pretraining for language understanding. Advances in neural information processing systems, 32, 2019.
Xinyu Ye, Hao Xiong, Jianhao Huang, Ziang Chen, Jia Wang, and Junchi Yan. An expressive quantum-driven graph learning approach with application to mixed-integer linear programming. In The Thirteenth International Conference Learning Representations, 2025.
14

---

Published as a conference paper at ICLR 2025
A DETAILED DERIVATION OF EQ. 6
To derive the expanded form of the KI. divergence in the learning objective, we start from the definition of the KL divergence, which is given
p(zr (2))
p(2:
1(2:-1/24 2,, Zo) (0 (z))
+ log Pa(zol qo Zo (z1)
CKL Bp( (m) log
Σ
log
(11)
((T))
1=2
Note that the KL. divergence between p and y is defined as
log P 4
CKL = Ep
(12)
Given that is p(Zorr] or) and q is go (zo:r (x); the KI. divergence should be
[oog p(zo:r go (2o:7 [x)
LKI. Eptzarlo)
(13)
Both and are Markov chains and can be factorized. The forward process p(Zo:r σ can be factorized
T
p(ze:r (or = (z (o) II [Zzllzz-1).
(14)
And the reverse process ge (zo.r (x) can also factorized as
T
go (x) = g(zr) II
to (2)-1
(15)
Here, Po(zolor is the initial distribution, (z,, 1) is the forward noising process, and go (z:-1 z, is the reverse denoising process. Substitute the factorized forms into the KL divergence and we can get
T
CKL Expza (
log pa (Z)lo) + Σ log p(zz(zz-
+ logg(Zz))
Σ
log go (z: (2 (2))))
T-1
=
=Eptzo )
(log08(20)) σr) log go 2n (2))
+
Σ
(log p (2:2222)))0) (1)22)1)
log (2: (2++))
+ log (oggp(z/2200 2T - log q(2T)).
(16)
We can express p(z/ p/zz/z(-11 in terms of p(zz-1/22 /2, 2o) using the properties of the Gaussian distributions in the diffusion process. First, recall that in the forward process, the joint distribution p(zo:r (or) can be rewritten using the posterior distributions p(zz-1/Z, 20):
T
p(zarT (or) = ps ( vo)
II p (z [e,, zo)-
(17)
This is possible because the diffusion process allows us to compute p(z;-1 Ze, 2o). which is Gaussian. Then substitute the new expression into the KL divergence:
LKL = Eppco.r (ar
log log (z (21) log p(zr|z)) (2o) - log g(Z)) +
(log p(zz 1//z--i/z,, zo) log go (z. 12:))
(18)
Now, the KL divergence becomes:
LKL Ep(zar ox) log (z(() + log Go (zo (2,
p(zz (Zo) (
p(z 22, Zu) 4o (2, (2,)
log
(19)
The derivation is completed. We can further interpret each term in the expanded KL divergence as follows.
15

---

Published as a conference paper at ICLR 2025
Initial KL Term (t 0).
log Pa 2o σ qo (o (1
E pp(
(20)
This term measures the discrepancy between the initial embedding distribution ps (zo or) and the model's 's initial reconstruction go (zolz)))
Final KL Term (t T).
log (T 20) ((zr)
(21)
This term compares the noise distribution at time T conditioned on Zo with the standard Gaussian prior (zr).
Intermediate KL Terms (2 StsT).
1 2o) log p(z 1 Ze, Zo)
(22)
(0 2, (2
These terms measure the discrepancy between the true posterior p(z: 1/4,-112,200 zo) and the model's predicted distribution ge 2 22).
B DETAILED DERIVATION OF EQ.
To derive the learning objective of Eq. 7, we first review some notations and definitions of diffusion models and the KL divergence in the main body the paper as follows.
(o1...., 0T.) sequence of discrete tokens. Physical Condition: a sequence physical variables conditioning the ground state of the Hamiltonian H(x). Token Embeddings: EMB( or) R lxxd continuous embeddings of the measurement string. Initial Embed- ding Distribution: pa [2 o = N (zo: EMB (r), Bo1). Forward Process: Adds noise to Zo to obtain 21, Zr, with transitions pl 2, (2:-)) N (2,; / 342--1 1,, BeT). Reverse Process: Starts from Zr (0, I) and predicts Zg using go Z 2 x) A (2 (Zz, x, t), T))
Measurement String:
Then derive this equation. The KL divergence between the forward process p p(Zo:rlo) and the reverse process go (Zo.7|Xx (x) is
CKL Eptao log p(zorrlo)
(23)
go (o(zo.Tl)
Expanding the joint distributions and we obtain
T II p(zilze-1)s
pl p(zorloo) = )o(2z)l))
t=1
(24)
go(zo.1(x) q(Zr) II
go 1z x2),
1=1
Substituting back into the KL divergence and obtain
LKL = (p(zoorlrrr
log (o((lo))) - log go (zo 21, x) +
(log [logpzz//zz11 - log go 10(2z-1/zz x)) log (
(25)
Then the KI. divergence is broken down and we can reorganize the KL. divergence into terms corresponding to each timestep. The formulation is given as
LKT. Ep(zo (ar) KL (v(zr(zo)]l((zr)) + Σ KL (p(z, Zo) (0)llga(22 x)) + KL (p(zo 0l0000/))))))
(26)
For Gaussian distributions, the KL divergence between p ~Npp Ep) and =N(4gg Eg) is
del En
det E,p )
kT((pl)4) = 2 tr( tr(E ) + (My - up) 1 (Mi 1p) - 1 + In
(27)
16

---

Published as a conference paper at ICLR 2025
where & is the dimensionality. Given that variances are fixed and known, and constants can be ignored during optimization, the KL divergence reduces to
KT. (pllg) 11 Hpll
(28)
Now we can specify the KL divergence each time step. For 1, the KL divergence between p(zo 21, σr) and qo x1, is:
KL (p(zo) Z1, σr) llu 21, x)) x ll (21, x, 1) (1 EMB(a))]I
(29)
where ji1 (Zz, EMB(or) )) is the true posterior mean from the forward process. When the time step t > 2, for each & we have
I| x, In the forward process, the posterior mean (rg linear combination of 2 and zo given by
KT (p(e-11 (2- x))
(
jie 20
(30)
(Z, zo) v zo
Vors(1
0x+-
Zz.
(31)
ar
Ch2
When Be is small, we can approximate jur (Z., zo) 2g- Substituting the approximations back into the KL divergence terms and we obtain the term for = 1 is
| (21, x, EMB (MMB(o))l (r)
(32)
The term for / 2 2 is
l| (Ze, x, 2ol1²
(33)
The rounding step involves mapping Zu back to the discrete measurements o using ğo (o(or)z)) 2o). The
negative ligelikelihood of this step is
-pp x.Z..7) log (o (2))
(34)
Then combine all terms, the learning objective becomes minimize
T
L = Ex,o
[|EBB(a) - (0(2), (0/2),,x. x, 1)11²
Σ
||z0 no x, 1)1 - log go (20)
(35)
This objective encourages the model to minimize the difference between the predicted means ne (,,, t) and the true values (EMB(o) at and Zo for t 2 2), and encourages to maxi- mize the likelihood of reconstructing the measurement string from
C MORE BASICS OF QUANTUM STATES AND QUANTUM MEASUREMENTS
In this section, we provide more details about the quantum states and quantum measurements and their relationship with classical joint distribution. For comprehensive discussion, refer the readers who are interested in quantum computing and quantum information to the Section 2.1 of the book Nielsen Chuang (2010).
A single qubit - the smallest unit of quantum computing is mathematically represented as a vector {55) = (0) + 311) parameterized by two complex numbers satisfying |a|= + [33 _ 1. Operations on a qubit must preserve this norm, and thus are described by 2 x umitary matrices. Of these, some of the most important are the Pauli operators; is useful to list them here:
0
x
Y
Z
(36)
-1
One could do some linear algebras and check that (0) [d] and [1)
[9] the eigenvectors of Z, |+) v | and 1-) = 1 [-]]] are the eigenvectors of X, (14) 1 and i vz 1
y
are the eigenvectors of Y. The same qubit can be decomposed into different orthonormal basis. For example,
(05) = a0) + B(1)
V2" (a B){+) + (a 3)
V2
(37)
V2 (a ( +
(a + (i) /2
17

---

Published as a conference paper at ICLR 2025
Positive-operne opperato valued measurement (POVM) is the testing manipulation of a physical sys- tem to yield numerical result. POVM is described by a set of measurement operators {T&,) ] satisfying Σx IIx = I and each I is positive semi-definite, where K is the total number of measurement operators. In this paper, we consider the Pauli-6 POVM (also named as random- ized single-qubit Pauli measurements in some literature) such that the measurement operators are (0) (0), 3 (1) (11, 31 --1>--++331 3
k=0
[ |}. It easy to check that these operators satisfy the POVM definition and K = 6. The reason for choosing the Pauli-6 POVM is that this measurement protocol is easy be implemented current quantum devices (NISQ devices) and is infornative-compllted (IC).
Measuring a qubit leads collapse of the qubit and produces an outcome with the probability p(k) satisfying the Born rule, which states that p(k) = tr(plIs), where = [z6 and (vb) is the transpose conjugate of ()). We may consider system of L qubits. It can be described by the wave function:
D Σ
Σ d==l
L V -........ (/) o1) OL);
where zz maps a fixed configuration o (01...., Sil...,..) σL) of qubits a complex number which is the amplitude satisfying Ex==1 Ex
( 0 1, and Ø; E (1.....kt K} is one of the K possible outcomes by performing quantum measurement the i-th qubit. It is formulated in a complex Hilbert space where the vector representation of the quantum state ( E Ck and its density matrix ( E x which becomes astronomical for large L.
Performing quantum measurement independently on qubits is easy be implemented. The most common strategy is to combine I. single-qubit measurement operators to I1,1 8 ..... 6 a where & the Kronecker product. Such measurement procedure outputs measurement Define p( . = (o).... oL)) 2 We can reformulate the wave function of quantum states to a clas- sical joint distribution. is a valid and legal joint distribution since Σ., Eos p(0 (00,,19))))
string (0I (o1...., OL) where 0; E {I....,K K} with probability | (0U.....,,,)).
E., Eo (o
and p(o1.....oL) 2
D DETAILS FOR BASELINE MODELS
We consider the classical shadow (CS) method (Huang et al., 2020), state-of-the-art, learning-free approach for efficiently constructing representations of unknown quantum states. For learning-basee comparisons, include kernel methods including the Radial Basis Function Kernel (RBFK) and Neural Tangent Kernel (NTK) following the implementation in Huang et al. (2022). Addition- ally, we evaluate advanced neural network-based models, including Recurrent Neural Network (RNN)-based architecture (Carrasquilla et al., 2019) and the transformer- based state-of of-the-art model LLM4QPE (Tang et al., 2024a). The details about their configurations are follows.
Classical Shadow (Huang et al., 2020). Classical Shadow (CS) is learning-free protocol used to efficiently predict many properties of quantum states with logarithmic number of measurements. By using randomized measurements and building memory-efficiee representation of the quantum state, it enables the estimation of various properties with high accuracy.
Kernel Methods (Huang et al., 2022). We utilize Radial Basis Function Kernel (RBFK) and Neural Tangent Kernel (NTK) to learn the feature map from physical conditions x to certain properties of the ground state of the Hamiltonian (x). A grid search is performed to identify the optimal regularization strength, with candidate values uniformly distributed on a logarithmic scale from 0.001 to 100. We employ a 5-fold cross-validatee strategy on the training dataset and present the predictive performance of the model that achieves the highest accuracy on the test dataset.
RNN (Carrasquilla et al., 2019). Recurrent Neural Networks (RNNs) are employed to reconstruct quantum states from measurement data, leveraging their ability to capture temporal or sequential dependencies. In our implementation, the RNN takes a sequence of measurement outcomes as input and learns to predict the underlying quantum state properties by modeling correlations between successive measurements. The model architecture includes a hidden layer with 128 units and is trained using the Adam optimizer. To ensure fair comparison, we perform early stopping based on validation loss and the test dataset for final evaluation of predictive performance.
18

---

Published as a conference paper at ICLR 2025
Table 3: RMSE of predicting the correlations of all subsystems of size two on the test dataset. The
result is averaged over Heisenberg model instances and each pair of adjacent qubits. For CS, M denotes the number of input measurements. While for the neural network-based approaches M denotes the number of sampled measurements Mout from trained models. The standard deviations are distinguished in gray.
L-
10
L
40
L- 70
L
100
Method
M
100
1000
10000
20000
100
1000
10000
20000
100
1000
10000
20000
100
1000
10000
20000
CS
0.1564
0,0009
00156
0/0107
0.1606
0.0638
0.0173
00121
[ 0.177
0.0545
00172
0.0121
0.1724
0,5477
0.0122
0.0122
RBFX
0.0796
0.0639
0.0578
0,0493
NTK
0.0755
0.0622
0.0865
0,0470
RNN
0.1328
0.044
0.0158
0,0502
0.0155
0.0034
0.0119
0.0014
0.1795
0.0388
0.0671
0.0132
0.0164
0.0035
0.016
00118
0.2137
0.0453
0,0739
0.0149
0.0240
0.0037
0.0153
0.0025
0.2525
0.0419
0.0006
0.0044
0.0251
0/0045
0.0163
0.0079
LLMKOPE
0.1316
0.0379
0.0140
0,0339
0.0032
0.8136
0,0993
0,0015
0.0392
0.1624
0,8513
0.0125
0,0142
0.0032
0.0017
0.0097
0.0462
0,1814
0,1527
0.0157
0.0155
0.0038
0.8116
0,0021
0,1739
0.0193
0J0531
0.0082
0,0152
0,0055
0,0114
0.0028
Outs
0,1269
0.0865
0,0433
0.0027
0.0022
0,0097
0,0085
0.0088
0,1582
0.0436
0.0115
0,9465
0.0113
0.0337
0.0034
0.0091
0.0479
0.3679
0.0138
0,0473
0.0117
0.0029
0.0011
0
0.1686
0.0512
0.0083
0.8478
0,0336
0,0125
0.0029
0,0098
Table 4: RMSE of predicting the entanglement entropies of all subsystems of size dataset. The result is averaged over Heisenberg model instances and each pair of adjacent qubits. For CS, MM denotes the number of input measurements. While for the neural network-based approaches M denotes the number of sampled measurements Mout from trained models. The standard deviations are distinguished in gray.
the test
1= 10
t= 30
5 = 70
L= 100)
Method
34 - 100
1000
10000
20000
100
1000
10000
20000
100
1000
10000
20000
100
1000
10000
20000
CS
0,5965
0,0922
0,0204
0,0119
0.6487
0.0927
0.0294
0.0259
0,6421
0,0943
0,0312
0,0298
0U6518
0.0998
0.0357
0,0316
RBEK
0.1268
1037
0.0977
0,0753
0.1379
0.0933
0.0719
0.1164
6.0187
0.0115
0.6132
0.1054
0U0346
0.0212
0.7948
0.3514
0.03K3
118299
0.1476
QJ1617
03439
010263
0,0022
0.0014
0.0314
0.001
0.0025
010857
0112000
LLM-OPE
0.4937
0.0948
0.0168
0.0176
0.0023
0.0102
0.5878
0.0A96
0U1223
3.0016
8.0307
0.001
0.6258
0.0912
0.8276
0.0016
0,0251
0.0088
8U6392
010979
0.1055
0,0322
0.0286
000018
0,0867
0.0132
0.0089
0,0861
0.0157
0.0145
0,0218
0,6125
0,0243
0.0019
0.0021
0.0014
LLM4QPE (Tang et al., 2024a). It presents a transformer-baset approach for predicting various properties of quantum systems through a pre- pre-training procedure that maximizes a likelihood function based on discrete measurement outcomes. This is also generative process similar to the conditional quantum state modeling discussed in Wang et al. (2022). The trained model can generate measure- ment samples conditioned the physical parameters which are unseen from the training dataset. Afterwards the quantum properties can be analyzed by using post- post-processing such as classical shadow. For fair comparison, we set the model with heads, 4 layers, and a hidden dimension of 128.
E VISUALIZED SUPPLEMENT OF THE COMPARISON BETWEEN QUADIM AND AUTOREGRESSIVE BASELINE
We include additional plot in Fig. 6, where present more fine-grained visualization of the prediction performance of QuaDiM. In this figure, each point represents the absolute error between the predieted correlation and the ground truth for all pairs of qubits across different Mout-
F ADDITIONAL EXPERIMENT RESULTS
F.1 IMPACT OF DIFFERENT PRE-DEFINED ORDER
To further compare the our model and the autoregressive baselines, we define an alternative ordering of qubits (from right to left or from the largest to the smallest index, as opposed to the left-to-right ordering used throughout the paper) and re- re-train the autoregressive baselines. The results of RMSE for predicting the correlations of all subsystems of size two the test dataset, evaluated in the predefined order from right to left, are presented in Tab. 5. For comparison, we also include experimental results already reported in the main text of the paper.
19

---

Published as a conference paper at ICLR 2025
Mour = 10000
QuaDIM LLMMQPE
Mase 1000
Maxe 100
tuobndirtd coteshetett dun
Mou
Mout 25
0.8
Absolute Error
Figure 6: Visualization of predicted correlations Cij for the ground state of the ID anti-ferromagnetic Heisenberg model of length L = 10 with different number of samples Mout from the trained models. We fix Myn 1000 and use classical shadow for proteesrcss..ng.. Each point represents the absolute error (lower is better) between the predicted correlation and the ground truth for all pairs of qubits across different Mott-
Table 5: RMSE of predicting the correlations of all subsystems of size two on the test dataset under different sequential order. M denotes the number of sampled measurements Mout from trained models. Notation <- denotes from right to left and - denotes from left to right.
Mothod
L= 70
100
M 100
M 1000)
M 10000
M 20000
M 100
M = 1000)
M 10000
M 20000
RNN(-)
0.2197
0.0165
0.0763
0.0159
RNNI >)
0.2137
0.0739
0.0240
0.0153
0.2325
0.0906
0.0251
0.0163
LLM-QPE(-)
0.1865
0,0538
0.0157
0,0108
0.1773
0,0542
0.0149
0,0122
LLMIOPE
0.1814
0.0531
QualiM
0.1679
0.9473
0.0117
0.0092
0.1686
0.0478
0.0125
0.0098
F.2 NUMERICAL RESULTS OF LEARNING LONG-RANGE XY MODEL
To further generalize the effectiveness of our proposed model, supplement experiments with results from a more physically general and classically challenging system: the long-range XY model in the presence of transverse field, whose Hamiltonian is given by
HXY = Σ (X, + Y,Y,) + Σ i<3
(38)
where Jij Jo/li - with 21. This quantum model inherits the long-range interactions between every two quantum sites, leading to a complex dynamics which is hard to simulated by classical computers. We restrict the system size L = 10 due to memory limitations. The ground states of quantum systems with different physical conditions Jij are calculated by eigenvalue decomposition. random sample series of Jı, and conduct elassical simulations to
Follow Xiao et al. (2022), collect the data. The experimental results for L = 10, with Mnn = 1000 and Mout E {1000, 10000}, are presented Tab. 6. It can be observed that QuaDiM consistently outperforms basclines.
Table 6: RMSE of predicting the correlations of ground states of the long-range XY model.
Method
Af = 10000
10010
CS
0.2575
0.0517
RBFK
0.1158
NTK
0.1939
0.2234
0.0502
LLMOQPE
0.2139
0.0482
0.0367
F.3 RESULTS OF VALIDATION ON OUT-OFFDIITTRBUUTION DATASET
Here, we consider evaluating the QuaDiM with out-of distriibuti (OOD) dataset, which means the dataset used for training and the dataset used for validation come from different distributions. We
20

---

Published as a conference paper at ICLR 2025
divide the sampled physical parameters x into two segments: the training set is limited to [0, 1.5]. while the test set exclusively spans [1.5, 2]. In alignment with setting of the paper, we set N = 100 and N = 20. We report the RMSE of predicted correlations for both the SOTA baseline LLM4QPE, and our model QuaDiM under both OOD and non-OOD (both training and fine-tuning are sampled from the same distribution [0, 2 conditions, with Mnn 1000 and Mout 10000. The results presented in the Tab. 7.
Table 7: RMSE of predicting the correlations of all subsystems of size two on the test dataset under the out-of-distribution (OOD) condition.
I. 70
I. 100
Method
no OOD
0OD
no OOD
OOD
M.M4
0.0155
0.0526
0.0152
0.0598 0.0125
QuaDiM
0.0117
0.0417
0.0465
F.4 MODEL SENSITIVITY TO THE DENOISING STEPS AND HIDDEN DIMENSION
In this section, we study the relationship between the hyper-parameters including the number of denoising steps Tr and the hidden dimension of QuaDiM and the prediction performance. We fix the number of diffusion steps during training for QuaDiM while shrinking the inference steps T using the approach introduced in DDIM (Song al., 2021a). We evaluate the model's performance under different inference step settings and compare it with both learning- free baseline, i.e. CS, and a learning- based SOTA model LLM4QPE. We report results for the task of predicting correlations with L = 100, M 1000, and Mout 1000.
As shown in the Tab. 8, when reducing inference to Tj 500 diffusion steps on a single GPU (2080T)), QuaDiM achieves a lower RMSE score compared to the CS while demonstrating an inference speed comparable to LLM4QPE.
To further investigate, we evaluate the model's performance the task of predicting correlations under a fixed dataset configuration L 10, Min = 1000, Mout = 1000 with different d values from (64, 128, 256, 512), and the resulting RMSE scores (lower is better) are: 0.0518, 0.0432, 0.0449, 0.0457, respectively. As the results show, setting = 128 achieves the best performance.
Table 8: RMSE of predicting correlations and the sampling speed with L, = 100, M., 1000, and Mout = 1000 under different denoising steps Ty.
Method
RMSE
Generated samples per sec.
CS
0.0547
LLM4QPE
0.0531
14.6
QuaDiM (T, 2000)
0.0478
5.7
QuaDiM (T, =1000)
0.0537
8.1
QuaDiM (T, -500)
0.0541
12.7
QuaDiM (T, 1000)
0.0882
37.4
F.5 MODEL SENSITIVITY TO THE POSITIONAL EMBEDDINGS
Positional embeddings in our context are employed capture the structural information among qubits. Additionally, we further investigate the model's performance in predicting correlations when using relative positional encoding (following the implementation in Shaw et (2018)) and no positional encoding at all. We report results for L 100 and M,,- 1000. The experimental results are provided Tab. 9. The results show that absolute and relative positional embeddings yield comparable performance, but removing positional information altogether significantly degrades the model's is ability to predict quantum correlations. This underscores the embeddings' role in preserving spatial relationships among qubits.
F.6 EXPERIMENTS ON TETRAHEDRAL POVM
In the main text of the paper, Pauli-6 POVM is used for data collection because this measurement protocol is easy implement current quantum devices and is informationally complete (IC). This
21

---

Published as a conference paper at ICLR 2025
Table 9: QuaDiM's RMSE of predicting correlations using different positional embedding (PE)
techniques.
Absolute PE
Relative PE
No PE
Mont 100
0.1686
0.1681
0.4527
Mont 0.0478
0.0482
0.3269
Mout 10000
0.0125
0.0139
0.2895
Mont 20000
0.0098
0.0110
0.2148
means that all the information of the quantum state can be recovered classically with a sufficiently large number of IC-POVM measurements. In other words, given the probability of each measurement outcome of IC-POVM, the quantum state can be uniquely determined.
To further validate method, here consider another type of IC-POVM: the tetrahedral POVM, to collect measurement data. The corresponding measurement operators are (I+ (()))(()))))))))) PP) (0,1,2.3)
where I the identity matrix, P represents the ensemble of Pauli operators (X, Z) and s
(0, O, 1).s(2) (2x2.0, - {)ss (- 12 3: 3 }).s(()) = (- --
- ---3.-33- that K = 4 for the tetrahedral POVM. We fix L = 10 with M,, = 1000, re-run the simulations to collect data, and re-train our model and the baselines. The numerical results are reported Tab. 10.
}. is easy to check
As shown, QuaDiM still outperforms the baselines in this scenario.
Table 10: RMSE of predicting correlations using tetrahedral POVM.
Method
M=1000
M=10000
CS
0.0512
0.0164
RBFK
0.0735
NTK
0.0747
RNN
0.0514
0.0163
LLM4QPE
0.0503
0.0141
QuaDiM
0.0433
0.0107
22


Reviews:
[{"review_id": "63CSf37IKp", "reviewer": "Reviewer_oe1d", "summary": "The paper proposes QuaDiM, a new framework based on diffusion models that learns the ground state of quantum systems conditioned on system parameters and enables quantum property estimation. QuaDiM uses a transformer-based architecture, where at inference time, measurement outcomes are first generated as continuous-valued hidden vectors and then decoded into discrete values. Numerical experiments are performed on 1D anti-ferromagnetic Heisenberg model up to $100$ qubits and predicted properties include correlation and entanglement entropy.", "questions": "* Can you discuss more on in what cases diffusion models would be strictly better than auto-regressive models, in the quantum state modeling setting?\n* I am not completely clear on how continuous latent vectors generated by QuaDiM are decoded back into measurement outcomes. Can you elaborate more on this?\n* Why QuaDiM will generate measurements consistently sampled from a valid quantum state without violating structure constraints. How is this enforced in the learning algorithm?\n* Does QuaDiM extrapolate to unseen system parameter $x$ distribution, likely with disjoint support as the training distribution, e.g., trained on $x \\in [0, 0.8]$, tested on $x \\in [0.9, 1]$?\n* Can you comment on QuaDiM's training efficiency? First, it seems that a relatively large $M_{in}$ is required for learning one ground state to good accuracy, also $M_{out}$ is not a small number as well. Can $M_{in}$ and $M_{out}$ be improved? Second, how many distinct pairs of (ground state, measurements) are needed for learning the mapping between Hamiltonian parameter and system ground state correctly, and how does this scale with the number of qubits?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper is very readable, with clear and understandable notations and easy-to-follow reasoning across the paper.\n* The proposed methods achieve SOTA performance on predicting correlation and entanglement entropy on 1D anti-ferromagnetic Heisenberg model with a reasonably large system size, justifying the effectiveness of the proposed approach.", "weaknesses": "* While the motivation for using non-autoregressive models is briefly discussed in the introduction, there could be more discussion and empirical evidence to support the claims about the sub-optimal performance of auto-regressive methods. While the measurements are not like language, the underlying physics model can still have some order structure (like a chain) and locality information is often important. Therefore, it's not completely obvious why diffusion models are necessarily better.\n* QuaDiM lacks machine learning technical novelty. The conditional diffusion model are already well studied in the literature and the used score neural network seems to be standard transformer as well. Given that ICLR is a ML conference, it would be more interesting to see applications and inspiration that can be drawn from QuaDiM to more general problems.\n* The experimental setup are not clearly presented. For example, the generation protocol of system parameter $x \\in \\mathbb{R}^{L - 1}$ is not described in the experiments section, not to mention the training and test set splitting process. This makes it hard to understand under what level QuaDiM extrapolates to \"unseen parameters\".\n* The experiment design is relatively weak. It would be interesting to consider more complicated property estimation problems, such as the estimation of entanglement negativity, two site observables with long ranges in between, or even the estimation of system parameters in the Hamiltonian. Second, it would be more convincing if QuaDiM is examined on a more complicated physics model, such as some with non-trivial long-range entanglement. It would be more convincing if the extrapolation of QuaDiM to critical state of the physics model could be tested and examined. \n* The benchmarks are not representative enough. For example, [1-3] are also works that model quantum states that should enable property estimation as downstream tasks, but are not discussed or compared.\n\nReferences:\n\n[1] Fitzek, David, et al. \"RydbergGPT.\" arXiv preprint arXiv:2405.21052 (2024).\n\n[2] Du, Yuxuan, et al. \"Shadownet for data-centric quantum system learning.\" arXiv preprint arXiv:2308.11290 (2023).\n\n[3] Zhang, Yuan-Hang, and Massimiliano Di Ventra. \"Transformer quantum state: A multipurpose model for quantum many-body problems.\" Physical Review B 107.7 (2023): 075147.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TeQ0qciuWa", "reviewer": "Reviewer_9ozo", "summary": "This paper introduces QuaDiM, a non-autoregressive diffusion model for quantum state property estimation (QPE). Unlike traditional autoregressive approaches that impose a (potentially biased) sequential ordering on qubits, QuaDiM treats all qubits equally through an iterative denoising process. The model is evaluated on the 1D anti-ferromagnetic Heisenberg model for systems up to 100 qubits, demonstrating superior performance in predicting correlation and entanglement entropy compared to baseline methods, especially with limited measurement data.", "questions": "1. How does the choice of diffusion schedule (β_t values) affect the model's performance? Would adaptive scheduling improve results?\n2. The paper uses Pauli-6 POVM measurements. How would the model perform with other measurement bases, and could this be made basis-independent?\n3. Could the authors comment on if the model can be used to obtain the log probabilities of each measurement outcome?\n4. POVM probabilities can be non-physical. Does the model always learn physical solutions?\n5. Could the model be modified to predict other quantum properties beyond correlation and entanglement entropy? What architectural changes would be needed?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed method of using diffusion model for quantum state tomography is novel. The paper effectively adapts diffusion models to quantum systems by introducing a token embedding function that maps discrete measurement outcomes to continuous features, allowing seamless integration with the diffusion process.\n2. The method seems to scale well to large systems. The model demonstrates strong performance on systems up to 100 qubits, which is significant for quantum computing applications.", "weaknesses": "1. While the paper mentions T=2000 denoising steps, there's no ablation study on how the number of steps affects performance versus computational cost.\n2. The paper fixes the embedding dimension at d=128 without analyzing how different dimensions affect the model's performance and computational requirements.\n3. The paper only considers Pauli-6 POVM, while evaluations on other POVM measurements may be beneficial.\n4. Although the authors made reasonable efforts, it is unclear to non-physicists how wave functions become probability distributions. Adding a section in the appendix discussing POVM measurements is preferable.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sNmcL0UA5X", "reviewer": "Reviewer_kHhj", "summary": "This paper addresses the challenge of Quantum state property estimation (QPE) in quantum many-body physics, focusing on predicting characteristics like correlation and entanglement entropy from measurement data. The authors introduce QuaDiM, a non-autoregressive generative model using diffusion models, which avoids the need for an intrinsic qubit ordering and treats all qubits equally and unbiasedly. QuaDiM learns to map physical variables to ground state properties during offline training and can sample from learned distributions for unseen variables to predict unknown quantum states. Empirical results on the 1D anti-ferromagnetic Heisenberg model with up to 100 qubits show that QuaDiM outperforms baseline models, particularly auto-regressive approaches, under limited training data and reduced inference sample complexity.", "questions": "Please correct me if I was wrong in terms of my previous comments :)", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 1, "strengths": "I am not an expert in quantum physics, so I can only provide my insights from the perspective of general generative models.\n\nFirst of all, the paper provides empirical evidence that models like diffusion models, which can capture the correlations between qubits, can outperform sequential-based models such as transformers and RNNs. This observation is an important finding in itself and can potentially steer research focus in the relevant AI4Science area.", "weaknesses": "As mentioned earlier, I would like to raise some concerns from a generative model perspective:\n\nFigure 3 is not informative to me; I can barely discern any difference between QuaDiM and LLM4QPE visually.\n\nBased on Figure 2, the methodological novelty appears limited.\n\nCould you please provide details about the computational resources used for the experiments? For example, the number of GPUs used, the training time, and the model size.\n\nIt would be beneficial if the authors could provide the standard deviation over seeds in Tables 1 and 2.\n\nMinor typo:\n\nLine 72: there is a double 'the'.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KNX63Y0Vat", "reviewer": "Reviewer_MPMV", "summary": "The paper proposes Quantum State Property Estimation using Diffusion Models (QuaDiM), a conditional diffusion model designed to learn and generate distributions of quantum states and accordingly quantum properties, for a given Hamiltonian parameter, from measurement data. Traditionally, auto-regressive models have been used for this task, modeling the quantum state array as a sequential structure despite the lack of inherent sequential ordering in quantum states. Unlike these approaches, QuaDiM denoises and shapes all states in the array simultaneously, as is typical in diffusion models, allowing it to effectively capture non-sequential interactions between spatially separated states. The proposed structure essentially includes an embedding that allows discrete quantum states to be treated as continuous diffusive evolution within a (conditional) diffusion model. The authors validate that the proposed QuaDiM outperformed existing main models, including RNN-based models and LMM4QPE, for large-scale quantum problems (e.g., L = 100).", "questions": "- Could the authors clarify why diffusion models have a suitable structure for learning non-sequential quantum state estimation compared to other generative models?\n\n- Could the authors explain why the embedding method used in QuaDiM has a structure suitable for non-sequential quantum state estimation?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed QuaDiM outperforms the baselines based on sequential models, e.g., RNN-based and transformer-based ones. It is intriguing, especially to handle very large-scale quantum states (e.g., L = 100) as a benchmark. \n\n- The paper is generally well-written.  I am not an expert in quantum computing, but I was able to follow this paper well and found it interesting to read.", "weaknesses": "- This paper is essentially an application paper that applies a diffusion model combined with a text embedding structure to the task of quantum state property estimation. It effectively addresses an interesting topic in the field of quantum computing through appropriate structural selection; however, it is uncertain whether this topic is broadly interesting to the ICLR audience.\n\n- I am not fully convinced why a diffusion model is the most suitable choice for the (non-sequential) quantum property estimation task. To me, this task fundamentally seems to be manageable by other generative models, such as conditional VAEs, GANs, or energy models. The authors mention that these generative models can only handle a single specific quantum state and cannot generalize to unseen states, but I am curious why a diffusion model is capable of generalizing to unseen states. While it is true that diffusion models demonstrate superior performance in terms of both fidelity and diversity, achieving high performance for diffusion models generally requires a large amount of training data. Therefore, the authors would clarify why the proposed QuaDiM would be particularly beneficial for quantum state property estimation, compared to other generative models, in a practical setting (e.g., limited observation as the authors mentioned).\n\n- The authors claim that the primary motivation behind the proposed QuaDiM model is to eliminate the sequential handling of quantum states present in existing models. However, QuaDiM employs an embedding function to represent discrete quantum states as continuous variables, following the approach in [1], which uses token and positional embeddings commonly used in text generation. It seems that this embedding, originally designed for text generation, inherently carries a sequential bias, which makes the motivation and effectiveness of the proposed QuaDiM less convincing to me.\n\n- The authors should explicitly mention the number of parameters and training time of the models used to ensure a fair comparison. Although this is not a critical issue, the network architecture and number of parameters of the main competitor, LLM4QPE, and the proposed QuaDiM appear to be nearly identical.\n\n- Some minor issues: In Table 2, for L = 70, M = 100 and QuaDiM (ours), is the result 0.0597 correct? Line 048, Fig. 2? Line 072, the the?, ...\n\n***\n[1] Li, X., Thickstun, J., Gulrajani, I., Liang, P. S., & Hashimoto, T. B. (2022). Diffusion-LM improves controllable text generation. Advances in Neural Information Processing Systems, 35, 4328-4343.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)