Quantum-JEPA: Quantum Born Priors for Latent-Action JEPA World Models
Organizations: Qinous Technology Co., Ltd., Beijing, China · Tsinghua University, Beijing, China · University of Notre Dame, Notre Dame, IN, USA
Abstract
A single observation can admit several plausible futures, making uncertainty an essential part of world modelling. However, quantum generative models remain underexplored for this task. We introduce Quantum-JEPA, a latent-action JEPA framework that learns distributions of future representations using a conditional quantum Born prior. Circuit measurements generate binary latent actions, which a shared predictor maps to alternative future representations. We train the quantum prior through the quality and diversity of these predictions using a sampled energy objective. Our model learns distinct continuations on constructed multimodal tasks and captures latent-action correlations that support coordinated multi-agent futures. On D3IL, a benchmark of diverse human demonstrations, our trained models generate multiple obstacle-crossing routes from a common starting point without observed motion. On the controlled tasks, our classical and quantum prior implementations improve prediction validity and distributional fidelity over the evaluated variational baselines. Our quantum priors combine competitive predictive quality with a parameter-efficient sampling structure: they avoid an exponentially sized categorical output layer while supporting correlated latent actions. We further validate finite-shot parameter-shift training, including training with readout noise. Unlike most prior work, which focuses on simulation, we also verify the feasibility of inference with trained priors on a real QPU. These results establish a trainable quantum component for JEPA future prediction and a basis for exploring quantum-enhanced world models at larger scales.
Figures & tables
| Prior | Parameters | Generation work |
|---|---|---|
| Softmax-Categorical | ||
| Independent-Bernoulli | ||
| Mixture- | ||
| AR-GRU | ||
| Conditional QCBM | classical quantum gates |
| Method / prior | Cov. | Valid | TV | minADE | ES |
|---|---|---|---|---|---|
| Variational JEPA | |||||
| Latent-CVAE | |||||
| Deterministic JEPA | |||||
| Uniform-Bernoulli | |||||
| Softmax-Categorical | |||||
| AR-GRU |
| Method / prior | Cov. | Valid | TV | minADE | ES |
|---|---|---|---|---|---|
| Variational JEPA | |||||
| Latent-CVAE | |||||
| Deterministic JEPA | |||||
| Uniform-Bernoulli | |||||
| Softmax-Categorical | |||||
| Independent-Bernoulli |
| Method / prior | Routes | Coverage | TV |
|---|---|---|---|
| Deterministic JEPA | – | ||
| Latent-CVAE | |||
| Softmax-Categorical | |||
| AR-GRU | |||
| Mixture-4 | |||
| Frozen-QCBM |
Appendix figures & tables16 assets
Supplementary material from the paper’s appendix.
Appendix
| Setting | MovingMNIST | HEP18 |
| Latent dimension / predictor width | 64 / 96 | 64 / 64 |
| Prior width | 48 | 48 |
| Bits / body layers | 3 / 2 | 12 / 4 |
| Terminal angles | 3 | 12 |
| World-model iterations | 4,000 | 6,000 |
| Batch size / prior samples | 64 / 32 | 32 / 32 |
| Method / prior | Prior | JEPA training | Readout | Deployment |
|---|---|---|---|---|
| Deterministic JEPA | 0 | 211,266 | 35,088 | 246,354 |
| Latent-CVAE | 4,296 | 169,248 | 35,088 | 184,728 |
| Softmax-Categorical | 15,792 | 227,058 | 35,088 | 262,146 |
| AR-GRU | 10,641 | 221,907 | 35,088 | 256,995 |
| Mixture-4 | 5,012 | 216,278 | 35,088 | 251,366 |
| Frozen-QCBM | 0 | 211,266 | 35,088 | 251,146 |
| Task / prior | Encoder | Predictor | Prior | Active | Readout |
|---|---|---|---|---|---|
| Digits / Deterministic JEPA | 439,600 | 50,112 | 0 | 0 | 12,966 |
| Digits / Uniform-Bernoulli | 439,600 | 50,112 | 0 | 0 | 12,966 |
| Digits / Frozen-QCBM | 439,600 | 50,112 | 3,435 | 3,423 | 12,966 |
| Digits / Softmax-Categorical | 439,600 | 50,112 | 3,640 | 3,640 | 12,966 |
| Digits / AR-GRU | 439,600 | 50,112 | 10,641 | 10,641 | 12,966 |
| Digits / Mixture-2 | 439,600 | 50,112 | 3,640 | 3,640 | 12,966 |
| Training | Cov. | Valid | TV | minADE | ES | Dice |
|---|---|---|---|---|---|---|
| Product-QCBM | ||||||
| Ideal gradient | ||||||
| Parameter shift | ||||||
| Shift + noise | ||||||
| Entangled-QCBM | ||||||
| Ideal gradient | ||||||
| Training | Cov. | Valid | TV | minADE | ES |
|---|---|---|---|---|---|
| Product-QCBM | |||||
| Ideal gradient | |||||
| Parameter shift | |||||
| Shift + noise | |||||
| Entangled-QCBM | |||||
| Ideal gradient | |||||
| Task | Method | Trainable parameters |
|---|---|---|
| MovingMNIST | Variational JEPA | 464,496 |
| MovingMNIST | Latent-CVAE | 472,208 |
| HEP18 | Variational JEPA | 22,288 |
| HEP18 | Latent-CVAE | 28,848 |
| Prior | Cov. | Valid | TV | |
|---|---|---|---|---|
| Entangled-QCBM | Yes | |||
| Entangled-QCBM | No | |||
| Mixture-4 | Yes | |||
| Mixture-4 | No |
| Prior | Setting | Coverage (%) | Validity (%) | TV |
|---|---|---|---|---|
| Entangled-QCBM | 82.60 | 98.59 | 0.31 | |
| Entangled-QCBM | 96.88 | 98.53 | 0.22 | |
| Entangled-QCBM | 99.83 | 98.55 | 0.16 | |
| Entangled-QCBM | 1 m / 0.5 m/s | 91.42 | 82.84 | 0.28 |
| Entangled-QCBM | 3 m / 1.5 m/s | 99.83 | 98.70 | 0.16 |
| Product-QCBM | 25.77 | 25.44 | 0.80 |
| Prior | Prior input | TV | Coverage | Valid |
|---|---|---|---|---|
| Product-QCBM | Original observation | |||
| Other image, same digit | ||||
| Other digit | ||||
| Training-input mixture | ||||
| Entangled-QCBM | Original observation | |||
| Other image, same digit |
| Prior | Sampling | Cov. | Valid | TV | minADE | ES | Dice |
|---|---|---|---|---|---|---|---|
| Product-QCBM | GPU / ideal | 89.06 | 98.44 | 0.1678 | 1.059 | 9.633 | 0.567 |
| Product-QCBM | QPU | 89.06 | 98.73 | 0.1674 | 1.059 | 9.670 | 0.567 |
| Entangled-QCBM | GPU / ideal | 96.88 | 96.48 | 0.1482 | 0.680 | 9.565 | 0.607 |
| Entangled-QCBM | QPU | 96.88 | 97.17 | 0.1520 | 0.680 | 9.548 | 0.607 |
| Sampling / layout | Cov. | Valid | TV | minADE | ES |
|---|---|---|---|---|---|
| Ideal sampling | 99.91 | 98.41 | 0.166 | 0.218 | 34.38 |
| QPU: original layout | 99.05 | 76.88 | 0.308 | 0.238 | 34.99 |
| QPU: selected layout | 98.78 | 81.53 | 0.280 | 0.241 | 34.86 |
| QPU: reflected layout | 99.57 | 86.69 | 0.239 | 0.227 | 34.66 |
| Prior | Sampling | Routes | TV |
|---|---|---|---|
| Product-QCBM | GPU (0% flips) | ||
| Product-QCBM | GPU (0.5% flips) | ||
| Product-QCBM | Real QPU | ||
| Entangled-QCBM | GPU (0% flips) | ||
| Entangled-QCBM | GPU (0.5% flips) | ||
| Entangled-QCBM | Real QPU |
| Prior | GPU (0%) | GPU (0.5%) | Real QPU |
|---|---|---|---|
| Product-QCBM | |||
| Entangled-QCBM |