Stochastic Schr\"odinger Diffusion Models for Pure-State Ensemble Generation

TL;DR

This paper introduces Stochastic Schr"odinger Diffusion Models (SSDMs), a novel score-based generative framework on the complex projective space for quantum pure-state ensemble generation, addressing geometric and intractability challenges.

Contribution

The paper develops SSDMs using Riemannian diffusion and a local Euclidean approximation, enabling training without explicit transition densities in quantum state space.

Findings

SSDMs accurately reproduce pure-state ensemble statistics

Generated quantum states improve quantum machine learning generalization

The method effectively captures observable moments, overlaps, and entanglement measures

Abstract

In quantum machine learning (QML), classical data are often encoded as quantum pure states and processed directly as quantum representations, motivating representation-level generative modeling that samples new quantum states from an underlying pure-state ensemble rather than re-preparing them from perturbed classical inputs. However, extending \emph{score-based} diffusion models with well-defined reverse-time samplers to quantum pure-state ensembles remains challenging, due to the non-Euclidean geometry of the complex projective space $\mathbb{CP}^{d-1}$ and the intractability of transition densities. We propose \emph{Stochastic Schr\"odinger Diffusion Models} (SSDMs), an intrinsic score-based generative framework on $\mathbb{CP}^{d-1}$ endowed with the Fubini--Study (FS) metric. SSDMs formulate a forward Riemannian diffusion with a stochastic Schr\"odinger equation (SSE) realization, and derive reverse-time dynamics driven by the Riemannian score $\nabla_{\mathrm{FS}} \log p_t$. To enable training without analytic transition densities, we introduce a local-time objective based on a local Euclidean Ornstein--Uhlenbeck approximation in FS normal coordinates, yielding an analytic teacher score mapped back to the manifold. Experiments show that SSDMs faithfully capture target pure-state ensemble statistics, including observable moments, overlap-kernel MMD, and entanglement measures, and that SSDM-generated quantum representations improve downstream QML generalization via representation-level data augmentation.