Quantum Oracle Sketching (QOS)

Code repository for our paper "Exponential quantum advantage in processing massive classical data".

In this work, we introduce Quantum Oracle Sketching, a framework that enables access to the classical world in quantum superposition for large-scale machine learning.

See the blog post at Quantum Frontiers for an introduction.

This repo contains:

• Core implementations of quantum oracle sketching in JAX.
• A tutorial that walk you through the basic usage of quantum oracle sketching.
• JAX implementation of quantum singular value transform (QSVT), including amplitude amplification, inversion, threshold, etc.
• A benchmark suite for quantum oracle and state sketching.
• Real-dataset experiments (classification and dimension reduction) for demonstrating the exponential memory advantage.

Overview

Quantum oracle sketching (QOS) is a quantum algorithm for loading classical data into a quantum computer. It instantiates the oracles needed by any quantum query algorithm using only random classical samples, with no full-dataset memory overhead.

This codebase includes two QOS simulation modes:

• qos_sampling.py: active random sampling implementation (more direct, heavier in simulation).
• qos.py: expected-unitary implementation (used for efficient benchmarking; conservative error upper bound).

They are implemented in JAX, which supports GPU/TPU execution and automatic differentiation.

Quick Start

1) Environment

python -m venv .venv
source .venv/bin/activate
pip install --upgrade pip
pip install -r requirements.txt

Notes:

• utils.py enables 64-bit JAX (jax_enable_x64=True).
• json is part of the Python standard library, so it is not listed in requirements.txt.
• If you want GPU JAX wheels, install JAX following official JAX instructions for your CUDA/TPU setup.

2) Tutorial: Interactive Notebook

If you want to play with QOS directly (instead of reproducing full paper figures), start with:

• notebooks/01_qos_quickstart.ipynb
  • Minimal, step-by-step examples for core QOS primitives.
  • Walkthrough for QOS on Boolean functions, general vectors, and matrix-element oracles.
  • Includes error-vs-samples scaling with fitted exponent.

Open Jupyter from repo root:

jupyter lab

If Jupyter is missing in your environment:

pip install jupyterlab ipykernel

3) Run the synthetic benchmark (main benchmark figure generation)

From repo root:

python benchmark.py

This generates:

• benchmark_flat_vector.pdf
• benchmark_general_vector.pdf
• benchmark_boolean_function.pdf
• benchmark_matrix_element.pdf
• benchmark_matrix_row_index.pdf

Repository Structure

.
├── benchmark.py              # Main benchmark used in the paper
├── data_generation.py        # Random sample generators for vectors/matrices/boolean functions
├── primitives.py             # Shared quantum primitives (e.g., amplitude amplification)
├── qos.py                    # QOS via expected unitaries (main benchmarking path)
├── qos_sampling.py           # QOS via explicit random sampling
├── qsvt.py                   # QSVT utilities + phase generation via pyqsp
├── utils.py                  # Numerical helpers, random instances, block-encoding helpers
├── notebooks/
│   └── 01_qos_quickstart.ipynb # Beginner tutorial notebook
├── real_datasets/            # Real-data experiments + plotting scripts
│   ├── *_svm.py              # LS-SVM-style classification accuracy vs machine size
│   ├── *_pca.py              # PCA variance recovery vs machine size
│   ├── *_combine_fig.py      # Combined 2-panel plots for each dataset
│   └── *_size_vs_*.json/.pdf # Outputs generated by running dataset scripts
└── requirements.txt

Real-Dataset Experiments

real_datasets/ evaluates machine-size vs performance under feature truncation and randomized feature sketches, for:

• IMDb sentiment (text TF-IDF)
• 20 Newsgroups topic data (text TF-IDF)
• PBMC68k single-cell RNA (UMI)
• Dorothea drug-discovery dataset
• Splice dataset (k-mer)

The mode-enabled real-dataset scripts use:

• --mode rare: rare-feature truncation (via min_df or min_samples).
• --mode bucket: balanced feature hashing / feature buckets.
• --mode jl: balanced signed sparse JL projection, implemented as the bucket transform with an additional random sign per original feature.

For bucket and jl, the requested sketch dimension is capped at the original feature dimension; this full-dimension endpoint returns the original data exactly. Bucket/JL random seeds are sampled reproducibly from a fixed sample seed in bucket_utils.py (default sample seed 42), with 5 seeds by default (--n-bucket-seeds or --n-jl-seeds).

Run full pipelines from scratch

Run from real_datasets/ so relative paths and defaults match script expectations.

IMDb:

python imdb_svm.py --mode rare
python imdb_pca.py --mode rare
python imdb_combine_fig.py --mode rare

20 Newsgroups (default averages over 100 random category pairs):

python 20news_svm.py --mode rare --n_pairs 100
python 20news_pca.py --mode rare --n_pairs 100
python 20news_combine_fig.py --mode rare

PBMC68k:

python pbmc68k_svm.py --mode rare
python pbmc68k_pca.py --mode rare
python pbmc68k_combine_fig.py --mode rare

Dorothea:

python dorothea_svm.py --mode rare
python dorothea_pca.py --mode rare
python dorothea_combine_fig.py --mode rare

Splice:

python splice_svm.py
python splice_pca.py
python splice_combine_fig.py

Bucket feature hashing and sparse JL numerics

Run from real_datasets/. A minimal single-dataset IMDb run for both SVM and PCA is:

python imdb_svm.py --mode bucket
python imdb_pca.py --mode bucket
python imdb_svm.py --mode jl
python imdb_pca.py --mode jl

After generating both bucket and JL JSON files for the same dataset, the combine scripts support --mode bucket_jl. This plots the bucket sparse / streaming / QOS curves and the sparse JL streaming curve in the same two-panel figure:

python imdb_combine_fig.py --mode bucket_jl

Dataset source/setup notes

• IMDb: auto-downloaded by imdb_utils.py from Stanford ACL IMDb.
• 20 Newsgroups: fetched through sklearn.datasets.fetch_20newsgroups.
• PBMC68k: loaded via scvelo.datasets.pbmc68k (downloaded/cached automatically).
• Dorothea: download manually from UCI (https://archive.ics.uci.edu/static/public/169/dorothea.zip) and extract to data_cache/dorothea (relative to where you run scripts).
• Splice: fetched via ucimlrepo (dataset id 69).

Core Files and Roles

• data_generation.py: sampling interfaces (vector_data, matrix_data, boolean_data).
• qos_sampling.py: explicit sampled-gate assembly; includes oracle/state tests under sampling.
• qos.py: expected-unitary assembly (used by benchmark.py).
• qsvt.py: polynomial angle generation + QSVT application helpers.
• primitives.py: amplitude amplification and related utilities.
• utils.py: random instance generators, block-encoding helpers, fidelity/infidelity utilities.

Reproducibility Notes

• Random seeds are fixed in scripts (np.random.seed(42) or JAX keys).
• Some full runs are heavy (large sample sweeps and repeated CV/SVD); expect long runtimes.
• Running qos.py or qos_sampling.py directly executes built-in tests with large default sizes.

Citation

If you find this repository useful, please consider citing our paper.

@article{zhao2026exponential,
    title={Exponential quantum advantage in processing massive classical data},
    author={Haimeng Zhao and Alexander Zlokapa and Hartmut Neven and Ryan Babbush and John Preskill and Jarrod R. McClean and Hsin-Yuan Huang},
    journal={},
    eprint={2604.07639},
    archivePrefix={arXiv},
    year={2026}
}