Why most quantum ML demos fall short
Quantum machine learning is often demonstrated through small, highly controlled examples: tiny datasets, shallow circuits, and weak baselines. While useful for education, these setups rarely convince practitioners.
Common issues include:
• Datasets too small to stress model capacity
• Comparisons against underpowered classical baselines
• High variance due to stochastic quantum training
• Little attention to reproducibility or deployment
Instead of building “a quantum model”, this project approaches quantum ML the same way we approach any new modeling family in applied ML: through a shared, reproducible experimental platform with strong baselines, realistic evaluation, and production considerations.
End-to-end platform overview
This project is not a single model but a config driven research and engineering framework for molecular property prediction.
It supports three model families under a unified pipeline:
• Classical graph neural networks (GNN, GAT, MPNN+GRU+Set2Set)
• Quantum models (variational circuits and quantum kernels)
• Hybrid classical quantum models with learned fusion
All models share the same:
• data ingestion
• featurization
• batching strategy
• training loop
• evaluation metrics
This design ensures that differences in performance can be attributed to architecture, not hidden changes in preprocessing or optimization
From SMILES to molecular graphs
Molecules are provided as SMILES strings, which are compact but poorly aligned with molecular interactions. To expose relational structure, each SMILES string is converted into a molecular graph using RDKit.
• Nodes represent atoms
• Edges represent bonds
• Directed edges enable message passing
Atom features encode chemically meaningful properties such as atom type, degree, hybridization, charge, and aromaticity. Bond features include bond type, conjugation, and ring membership.
Why geometry matters (especially for QM9)
Many molecular properties depend not only on connectivity, but also on three-dimensional geometry. This is particularly true for the electric dipole moment μ, which depends on the spatial distribution of partial charges.
To expose this information, the platform optionally augments molecular graphs with 3D features. Hydrogen atoms are added, a conformer is generated using RDKit’s ETKDG algorithm, and the structure is relaxed via MMFF optimization.
Each atom receives Cartesian coordinates and its distance to the molecular centroid, while edges are augmented with bond length information. Importantly, the same featurization pipeline feeds all model families, isolating architectural effects rather than input differences.
Scaling graph learning in JAX
Graph neural networks introduce a practical challenge: molecules have variable size. This clashes with JAX’s preference for static tensor shapes.
To address this, the platform introduces a BatchedGraph abstraction:
• Node and edge features are padded to batch level maxima
• Boolean masks track real vs padded entries
• All aggregation and pooling operations are mask-aware
This enables:
• true mini-batch training
• JIT compiled forward passes
• stable gradient statistics
• efficient hyperparameter sweeps
In practice, this approach yields roughly an order of magnitude speedup over naive per graph execution.
Three model families under one roof
Classical baselines that actually matter
Any meaningful evaluation of quantum or hybrid models requires strong classical baselines.
The platform includes:
• a message passing GNN baseline
• a Graph Attention Network (GAT) with edge aware attention
• an MPNN with GRU updates and Set2Set pooling
All classical models share the same training and evaluation pipeline, ensuring fair comparison.
Quantum models beyond toy circuits
A legacy 4-qubit variational circuit is included as a reference point, establishing a lower bound on quantum capacity.
More expressive quantum models are then explored:
• 8-qubit variational circuits
• data re-uploading feature maps
• hardware-aware circuit design
• quantum kernel methods as a non-variational alternative
Noise, shot statistics, and gradient behavior are treated as first-class concerns rather than ignored artifacts.
Hybrid classical–quantum models done properly
Naive hybrid models often concatenate small quantum outputs with large classical embeddings, causing the quantum signal to be ignored.
To address this, the Hybrid V2 architecture enforces dimension-matched fusion:
• a shared GNN encoder produces a 128-D embedding
• a classical branch projects to 64-D
• a quantum branch compresses to qubits, processes via a VQC, then upsamples back to 64-D
• a learned gate combines both representations per input
Hybrid training is staged:
1. classical pretraining
2. joint training
3. low-learning-rate fine-tuning
This allows quantum features to contribute conditionally, not by assumption.
Training, evaluation, and scientific rigor
Experiments are defined through configuration files and hashed for reproducibility. Random seeds are controlled across Python, NumPy, and JAX.
Evaluation includes:
• Murcko scaffold splits for realistic generalization
• systematic ablations over architecture and circuit parameters
• uncertainty estimation via MC dropout and quantum shot noise
• diagnostics for barren plateaus and trainability
• surrogatability tests to probe classical mimicry of quantum outputs
The goal is not headline numbers, but understanding model behavior.
From research code to production inference
Trained models are managed via a lightweight model registry that handles configuration, checkpoint loading, and caching.
A FastAPI service exposes:
• single-molecule inference
• batch inference
• optional uncertainty estimates
From the client’s perspective, classical, quantum, and hybrid models share the same interface. This closes the gap between experimentation and real-world use.
Code and reproducibility
The full implementation — including configuration files, data loaders, model architectures, training pipelines, evaluation utilities, and the FastAPI inference service — is available here:
GitHub repository: https://github.com/Pratik25priyanshu20/Quantum-ML-Drug-discovery
Final thoughts
Rather than asking whether quantum models “win”, this platform asks a more useful question: when, where, and how do they contribute?
Answering that responsibly requires shared infrastructure, strong baselines, and measurable uncertainty. This project aims to provide exactly that.
Top comments (0)