diff --git a/EQNN_Cosmos_Dong/Notebooks/EQNN_SO(2).ipynb b/EQNN_Cosmos_Dong/Notebooks/EQNN_SO(2).ipynb index 16e1036..565d31e 100644 --- a/EQNN_Cosmos_Dong/Notebooks/EQNN_SO(2).ipynb +++ b/EQNN_Cosmos_Dong/Notebooks/EQNN_SO(2).ipynb @@ -1,5 +1,18 @@ { "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "1770905f", + "metadata": {}, + "outputs": [], + "source": [ + "# Install dependencies for Google Colab/Codespaces\n", + "import sys\n", + "if 'google.colab' in sys.modules or 'vscode' in sys.modules:\n", + " !pip install -q jax flax pennylane optax" + ] + }, { "cell_type": "code", "execution_count": 1, diff --git a/EQNN_Cosmos_Dong/Notebooks/EQNN_discrete.ipynb b/EQNN_Cosmos_Dong/Notebooks/EQNN_discrete.ipynb index 5b07e01..4671e75 100644 --- a/EQNN_Cosmos_Dong/Notebooks/EQNN_discrete.ipynb +++ b/EQNN_Cosmos_Dong/Notebooks/EQNN_discrete.ipynb @@ -1,5 +1,18 @@ { "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "67f284cb", + "metadata": {}, + "outputs": [], + "source": [ + "# Install dependencies for Google Colab/Codespaces\n", + "import sys\n", + "if 'google.colab' in sys.modules or 'vscode' in sys.modules:\n", + " !pip install -q jax flax pennylane optax" + ] + }, { "cell_type": "code", "execution_count": 1, @@ -23,11 +36,7 @@ "outputs": [ { "data": { - "application/javascript": [ - "MathJax.Hub.Config({\n", - " TeX: {equationNumbers: { autoNumber: \"AMS\" } }\n", - "});\n" - ], + "application/javascript": "MathJax.Hub.Config({\n TeX: {equationNumbers: { autoNumber: \"AMS\" } }\n});\n", "text/plain": [ "" ] diff --git a/EQNN_Cosmos_Dong/Notebooks/EQNN_discrete_Asymm.ipynb b/EQNN_Cosmos_Dong/Notebooks/EQNN_discrete_Asymm.ipynb index c689986..3d3f603 100644 --- a/EQNN_Cosmos_Dong/Notebooks/EQNN_discrete_Asymm.ipynb +++ b/EQNN_Cosmos_Dong/Notebooks/EQNN_discrete_Asymm.ipynb @@ -1,5 +1,18 @@ { "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "82916fe9", + "metadata": {}, + "outputs": [], + "source": [ + "# Install dependencies for Google Colab/Codespaces\n", + "import sys\n", + "if 'google.colab' in sys.modules or 'vscode' in sys.modules:\n", + " !pip install -q jax flax pennylane optax" + ] + }, { "cell_type": "code", "execution_count": 1, diff --git a/EQNN_Cosmos_Dong/Notebooks/SU(2)_trainable.ipynb b/EQNN_Cosmos_Dong/Notebooks/SU(2)_trainable.ipynb index e1f7308..40f3072 100644 --- a/EQNN_Cosmos_Dong/Notebooks/SU(2)_trainable.ipynb +++ b/EQNN_Cosmos_Dong/Notebooks/SU(2)_trainable.ipynb @@ -1,5 +1,18 @@ { "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "3d9c03a3", + "metadata": {}, + "outputs": [], + "source": [ + "# Install dependencies for Google Colab/Codespaces\n", + "import sys\n", + "if 'google.colab' in sys.modules or 'vscode' in sys.modules:\n", + " !pip install -q jax flax pennylane optax" + ] + }, { "cell_type": "code", "execution_count": 1, diff --git a/EQNN_Cosmos_Dong/Notebooks/SU(2)_untrainable.ipynb b/EQNN_Cosmos_Dong/Notebooks/SU(2)_untrainable.ipynb index da0c53c..d26ef69 100644 --- a/EQNN_Cosmos_Dong/Notebooks/SU(2)_untrainable.ipynb +++ b/EQNN_Cosmos_Dong/Notebooks/SU(2)_untrainable.ipynb @@ -1,5 +1,18 @@ { "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "577a81df", + "metadata": {}, + "outputs": [], + "source": [ + "# Install dependencies for Google Colab/Codespaces\n", + "import sys\n", + "if 'google.colab' in sys.modules or 'vscode' in sys.modules:\n", + " !pip install -q jax flax pennylane optax" + ] + }, { "cell_type": "code", "execution_count": 1, diff --git a/EQNN_Cosmos_Dong/requirements.txt b/EQNN_Cosmos_Dong/requirements.txt new file mode 100644 index 0000000..f766cdd --- /dev/null +++ b/EQNN_Cosmos_Dong/requirements.txt @@ -0,0 +1,4 @@ +jax>=0.9.0.1 +flax>=0.12.4 +pennylane>=0.44.0 +optax>=0.1.0 diff --git a/Quantum_AE_for_HEP_Tom_Magorsch/dev_notebooks/Electron-Photon/README.md b/Quantum_AE_for_HEP_Tom_Magorsch/dev_notebooks/Electron-Photon/README.md new file mode 100644 index 0000000..402b700 --- /dev/null +++ b/Quantum_AE_for_HEP_Tom_Magorsch/dev_notebooks/Electron-Photon/README.md @@ -0,0 +1,70 @@ +# Quantum & Classical Autoencoders for Electron-Photon Classification + +This folder contains a collection of notebooks exploring **autoencoder architectures** (both classical and quantum-enhanced) for distinguishing electrons from photons in high-energy physics calorimeter data. + +## Physics Objective + +Electron-Photon discrimination is a critical classification task in particle physics. Both particles leave similar electromagnetic (EM) signatures in calorimeters, making their separation challenging. This folder explores how classical autoencoders (CAE, VAE) and quantum-enhanced autoencoders (QAE, SQAE) can learn discriminative representations of calorimeter shower patterns for robust electron-photon tagging. + +## Dataset + +- **LHC Electron-Photon Dataset**: Calorimeter images from simulated electron and photon events. +- **File Format**: HDF5 (.hdf5) containing EM calorimeter shower images (X) and binary labels (y: 0=photon, 1=electron). +- **Typical Shapes**: Multi-thousand samples of 2D/3D calorimeter energy deposits. + +## Notebooks Overview + +### Classical Baselines +- **Electron-Photon_CAE_fulldata.ipynb**: Convolutional AutoEncoder for full dataset feature learning. +- **Electron-Photon_CAE_keras_tuner.ipynb**: Hyperparameter-tuned CAE using Keras Tuner. +- **Electron-Photon_VAE_fulldata.ipynb**: Variational AutoEncoder for probabilistic shower representation. + +### Quantum-Enhanced Models +- **Electron-Photon_QAE.ipynb**: Quantum AutoEncoder with quantum latent space. +- **Electron-Photon_QAE_fulldata.ipynb**: Full-dataset variant of QAE. +- **Electron-Photon_SQAE.ipynb**: Single-qubit Quantum AutoEncoder. +- **Electron-Photon_SQAE_fulldata.ipynb**: Full-dataset training variant. + +### Data Utilities +- **EMD_experiment.ipynb**: Earth Mover Distance (EMD) analysis of shower patternsfor model comparison. + +### Advanced Models +- **gammaetune/**: Hyperparameter optimization workflows for gamma/electron discrimination. + +## Key Results + +- Classical autoencoders (CAE/VAE) achieve competitive classification metrics on EM calorimeter data. +- Quantum autoencoders show promise in learning compressed shower representations. +- Full-data training variants improve generalization compared to subset-trained models. + +## Dependencies + +Install via the parent directory's `requirements.txt`: + +```bash +pip install -r ../requirements.txt +``` + +Key packages: +- `tensorflow` / `tensorflow-quantum` +- `pennylane` for quantum circuits +- `h5py` for data loading +- `scikit-learn` for preprocessing +- `keras-tuner` for hyperparameter optimization + +## Usage + +1. **Download data**: Place HDF5 files (`electron.hdf5`, `photon.hdf5`) in `../../data/` (adjust paths as needed). +2. **Run a notebook**: Open any `.ipynb` file in Jupyter and execute cells from top to bottom. +3. **Compare models**: Each architecture can be independently evaluated; check training curves and ROC/AUC metrics. + +## Future Directions + +- Scale to full calorimeter geometries (e.g., 100×100 pixel grids). +- Implement quantum error mitigation for accurate circuit simulation. +- Explore attention-based variants (ViT-style) for multi-scale feature learning. +- Extend to JAX/Flax for improved computational efficiency. + +--- + +**Contributing**: This folder is part of the ML4SCI/QMLHEP initiative to advance quantum machine learning in high-energy physics research. diff --git a/Quantum_AE_for_HEP_Tom_Magorsch/dev_notebooks/MNIST/README.md b/Quantum_AE_for_HEP_Tom_Magorsch/dev_notebooks/MNIST/README.md new file mode 100644 index 0000000..d3c0afe --- /dev/null +++ b/Quantum_AE_for_HEP_Tom_Magorsch/dev_notebooks/MNIST/README.md @@ -0,0 +1,73 @@ +# Quantum & Classical Autoencoders on MNIST + +This folder contains a comprehensive benchmark suite of **autoencoder architectures** (both classical and quantum-enhanced) evaluated on the MNIST handwritten digit dataset. + +## Objective + +The MNIST dataset serves as a **controlled benchmark** for developing and validating autoencoder designs before applying them to high-energy physics data. This folder explores how quantum autoencoders compare to classical baselines in terms of compression efficiency, reconstruction fidelity, and trainability on a well-understood task. + +## Dataset + +- **MNIST**: 28×28 grayscale images of handwritten digits (0–9). +- **Source**: Built-in TensorFlow/Keras dataset. +- **Typical Shapes**: 60k training samples, 10k test samples, normalized to [0, 1]. + +## Notebooks Overview + +### Classical Baselines +- **MNIST_CAE.ipynb**: Standard Convolutional AutoEncoder for image compression. +- **MNIST_convolutional_CAE.ipynb**: Advanced convolutional variant with detailed architecture. +- **MNIST_PCA_CAE.ipynb**: Hybrid approach combining PCA with convolutional layers. + +### Quantum-Enhanced Models +- **MNIST_QAE.ipynb**: Quantum AutoEncoder with quantum latent space for digit representation. +- **MNIST_SQAE.ipynb**: Single-Qubit Quantum AutoEncoder variant for minimal quantum resource usage. + +### Hyperparameter Optimization +- **MNIST_QAE_hyperparameters.ipynb**: Systematic tuning of quantum circuit depth, learning rates, and layer configurations. +- **MNIST_QAE_model.ipynb**: Best-performing QAE configuration discovered through tuning. +- **MNIST_SQAE_hyperparameters.ipynb**: Hyperparameter search for single-qubit variant. + +## Key Metrics + +- **Reconstruction MSE**: Measures how well the autoencoder reconstructs digits. +- **Latent Dimensionality**: Compression ratio achieved (original size vs. bottleneck). +- **Training Efficiency**: Wall-clock time and convergence speed for classical vs. quantum variants. + +## Typical Results + +- Classical CAE: MSE ~0.01–0.02 on test set; ~10× compression ratio. +- Quantum AE: Competitive reconstruction with fewer latent qubits; scalability limited by circuit depth. +- Hyperparameter tuning improves both classical and quantum variants by ~5–15%. + +## Dependencies + +Install via the parent directory's `requirements.txt`: + +```bash +pip install -r ../requirements.txt +``` + +Key packages: +- `tensorflow` / `tensorflow-quantum` +- `pennylane` for quantum circuits +- `cirq` for quantum circuit visualization +- `scikit-learn` for PCA and preprocessing + +## Usage + +1. **Run a notebook**: Open any `.ipynb` file in Jupyter/Colab and execute cells from top to bottom. +2. **Compare architectures**: Each notebook produces training curves and reconstruction visualizations. +3. **Tune hyperparameters**: Start with `*_hyperparameters.ipynb` notebooks to explore the parameter space. +4. **Visualize results**: Check digit reconstruction plots and latent space embeddings. + +## Future Directions + +- Extend quantum circuits to deeper, more expressive topologies. +- Implement variational quantum autoencoders (VQAE) with Barren Plateau mitigation. +- Port to JAX/Flax for improved numerical stability and GPU acceleration. +- Apply learned representations as pretrained features for downstream HEP tasks (Quark-Gluon, Electron-Photon). + +--- + +**Contributing**: This folder is part of the ML4SCI/QMLHEP initiative advancing quantum machine learning benchmarks and techniques. diff --git a/Quantum_AE_for_HEP_Tom_Magorsch/dev_notebooks/Pennylane/README.md b/Quantum_AE_for_HEP_Tom_Magorsch/dev_notebooks/Pennylane/README.md new file mode 100644 index 0000000..18c29fe --- /dev/null +++ b/Quantum_AE_for_HEP_Tom_Magorsch/dev_notebooks/Pennylane/README.md @@ -0,0 +1,87 @@ +# PennyLane Quantum Autoencoders for HEP & ML + +This folder contains a comprehensive exploration of **quantum autoencoder designs** implemented using **PennyLane**, including hybrid classical-quantum architectures and advanced optimization techniques. + +## Objective + +These notebooks showcase the evolution and optimization of quantum autoencoders (QAE, SQAE) for high-energy physics tasks and general machine learning benchmarks. A key focus is comparing different optimization strategies (gradient descent, QNG, parameter initialization) to improve convergence and reconstruction quality. + +## Architecture Families + +### Hybrid Autoencoders (HAE) +- **pennylane-HAE.ipynb**: Basic Hybrid AutoEncoder combining classical CNN encoder with quantum bottleneck. +- **pennylane-HAE-QG.ipynb**: HAE tuned for Quark-Gluon jet data. +- **pennylane-HAE (2) (2).ipynb**: Variant with improved circuit depth control. + +### Single-Qubit Quantum Autoencoders (SQAE) +- **pennylane-SQAE.ipynb**: Fundamental SQAE architecture. +- **pennylane-SQAE-keras.ipynb**: SQAE integrated with Keras/TensorFlow training. +- **pennylane-hybrid.ipynb**: Hybrid training combining classical and quantum layers. +- **pennylane-hybridSQAE.ipynb**: Full hybrid workflow. + +### Advanced Optimization Variants +- **pennylane-SQAE-QNG.ipynb**: Quantum Natural Gradient (QNG) optimizer for improved convergence. +- **pennylane-SQAE-ampenc.ipynb**: Amplitude encoding for efficient state preparation. +- **pennylane-SQAE-DRC.ipynb**: Dynamical Reconfigurable Circuits. +- **QNG_comparasion.ipynb**: Benchmarks QNG vs. standard gradient descent. + +### Physics-Specific Models +- **pennylane-SQAE-convDRC-QG.ipynb**: Convolutional DRC for Quark-Gluon classification. +- **pennylane-SQAE-convDRC-QG-jax.ipynb**: Same architecture in JAX/Flax backend. +- **pennylane-SQAE-convDRC-QG-model.ipynb**: Production-ready configuration. +- **pennylane-SQAE-convDRC-EP.ipynb**: Electron-Photon variant. +- **pennylane-SQAE-convDRC-Qconv.ipynb**: Quantum convolutional layers. + +### Supervised Learning with Quantum Circuits +- **pennylane-supervised-qg.ipynb**: Quantum circuit for supervised Quark-Gluon classification. +- **pennylane-supervised-qg-Copy1.ipynb**: Variant of above. + +### Debugging & Development +- **Batched_hybrid_debug.ipynb**: Troubleshooting batched quantum-classical hybrid training. +- **animation.gif, animation2.gif, animation3.gif**: Training visualizations (Bloch sphere rotations, loss curves). +- **bloch.gif, bloch2.gif, bloch3.gif**: Animated Bloch sphere representations of learned quantum states. + +## Key Innovations + +- **Quantum Natural Gradient (QNG)**: Faster convergence than vanilla gradient descent (see `QNG_comparasion.ipynb`). +- **Dynamical Reconfigurable Circuits**: Adapting quantum circuit structure during training. +- **JAX Backend**: Several notebooks use JAX for automatic differentiation and GPU acceleration. +- **Batch Training**: Hybrid architectures supporting mini-batch updates (see `Batched_hybrid_debug.ipynb`). + +## Typical Performance + +- Hybrid Autoencoders: ~0.75–0.85 AUC on Quark-Gluon; ~50% latent dimension reduction. +- SQAE Models: Scalable to 200+ layers with careful initialization and learning rate tuning. +- QNG Optimization: 2–5× faster convergence vs. standard SGD on small circuits. + +## Dependencies + +Install via the parent directory's `requirements.txt`: + +```bash +pip install -r ../requirements.txt +``` + +Essential packages: +- `pennylane` (quantum circuits & autoencoders) +- `pennylane-qiskit` or `pennylane-lightning` (simulator backends) +- `jax` / `jaxlib` (for JAX-based notebooks) +- `tensorflow` / `tensorflow-quantum` (for Keras integration) + +## Usage + +1. **Start with basics**: Run `pennylane-HAE.ipynb` or `pennylane-SQAE.ipynb` for fundamentals. +2. **Compare optimizers**: Use `QNG_comparasion.ipynb` to understand QNG benefits. +3. **Physics applications**: Jump to `pennylane-SQAE-convDRC-QG*.ipynb` for Quark-Gluon tasks. +4. **Visualize training**: Check `.gif` files for animated Bloch sphere and loss curves. + +## Future Directions + +- Port all models to full JAX/Flax ecosystem for seamless integration with QMLHEP7. +- Implement error mitigation techniques for hardware execution. +- Extend to multi-qubit entangling circuits for increased expressivity. +- Auto-tune hyperparameters using Bayesian optimization or RL. + +--- + +**Contributing**: These notebooks form a foundation for quantum autoencoder research in ML4SCI/QMLHEP, enabling rapid prototyping and physics-aware optimization. diff --git a/Quantum_AE_for_HEP_Tom_Magorsch/dev_notebooks/Quark-Gluon/README.md b/Quantum_AE_for_HEP_Tom_Magorsch/dev_notebooks/Quark-Gluon/README.md new file mode 100644 index 0000000..8382c59 --- /dev/null +++ b/Quantum_AE_for_HEP_Tom_Magorsch/dev_notebooks/Quark-Gluon/README.md @@ -0,0 +1,70 @@ +# Quantum & Classical Autoencoders for Quark-Gluon Jet Classification + +This folder contains a collection of notebooks exploring **autoencoder architectures** (both classical and quantum-enhanced) for distinguishing quark-initiated jets from gluon-initiated jets in high-energy physics. + +## Physics Objective + +Quark-Gluon jet discrimination is a fundamental classification task in particle physics. Distinguishing quarks from gluons helps identify the underlying hard process in collider experiments. This folder explores how classical autoencoders (CAE, VAE, DAE) and quantum-enhanced autoencoders (QAE, SQAE) can learn compressed representations of jet structures for improved classification. + +## Dataset + +- **LHC Quark-Gluon Jet Dataset**: Multi-dimensional jet feature data from simulated particle collisions. +- **File Format**: HDF5 (.hdf5) containing jet features (X) and binary labels (y: 0=gluon, 1=quark). +- **Typical Shapes**: Training sets of ~793k samples, test sets of 10k–139k samples with 40×40 spatial jet representations. + +## Notebooks Overview + +### Classical Baselines +- **Quark-Gluon-CAE.ipynb**: Convolutional AutoEncoder baseline for unsupervised feature learning. +- **Quark-Gluon-VAE.ipynb**: Variational AutoEncoder for probabilistic jet representation. +- **Quark-Gluon-DAE.ipynb**: Denoising AutoEncoder for robust feature extraction. + +### Quantum-Enhanced Models +- **Quark-Gluon-QAE.ipynb**: Quantum AutoEncoder combining classical encoders with quantum latent layers. +- **Quark-Gluon-SQAE.ipynb**: Single-qubit Quantum AutoEncoder variant. +- **ViT-QSAL.ipynb**: Vision Transformer with Quantum Self-Attention Layers for attention-based jet classification. + +### Data Utilities +- **Explore_Data.ipynb**: EDA of jet structure, feature distributions, and class balance. +- **Quark-Gluon-data-scaling.ipynb**: Normalization and preprocessing pipelines. +- **Compress_Data.ipynb**: Methods for reducing data size for efficient training. + +### Advanced Models +- **Coatnet.ipynb**: Hybrid convolutional-attention architecture for jet classification. +- **swae.ipynb**: Sliced Wasserstein AutoEncoder variant. + +## Key Results + +- Classical autoencoders (CAE/VAE) achieve ~0.75–0.80 AUC on test sets. +- Quantum-enhanced variants show competitive or improved performance with reduced latent dimensionality. +- Vision Transformer + Quantum Attention demonstrates the feasibility of attention-based quantum circuits for jet data. + +## Dependencies + +Install via the parent directory's `requirements.txt`: + +```bash +pip install -r ../requirements.txt +``` + +Key packages: +- `tensorflow` / `tensorflow-quantum` +- `pennylane` for quantum circuits +- `h5py` for data loading +- `scikit-learn` for preprocessing + +## Usage + +1. **Download data**: Place HDF5 files in `../../data/quark-gluon/` (adjust paths as needed). +2. **Run a notebook**: Open any `.ipynb` file in Jupyter and execute cells from top to bottom. +3. **Compare models**: Each notebook can be run independently; visualizations compare reconstruction quality and classification metrics. + +## Future Directions + +- Extend to full JAX/Flax ecosystem for improved scalability. +- Implement distributed training for larger datasets. +- Explore error mitigation techniques for NISQ hardware simulation. + +--- + +**Contributing**: This folder is part of the broader ML4SCI/QMLHEP effort to advance quantum machine learning in high-energy physics research. diff --git a/Quantum_Attention_Benchmarks/JAX_Quantum_Attention_Baseline.ipynb b/Quantum_Attention_Benchmarks/JAX_Quantum_Attention_Baseline.ipynb new file mode 100644 index 0000000..a867f61 --- /dev/null +++ b/Quantum_Attention_Benchmarks/JAX_Quantum_Attention_Baseline.ipynb @@ -0,0 +1,151 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "82da31e7", + "metadata": {}, + "source": [ + "# JAX + Flax Quantum Attention Baseline\n", + "\n", + "This notebook demonstrates a minimal integration of a **multi-head attention**-style\n", + "module built with Flax and a **PennyLane FlaxLayer** to show compatibility with the\n", + "JAX/Flax ecosystem. It serves as a toy benchmark for the QMLHEP7 quantum particle\n", + "transformer project. The classical attention block is replaced by a quantum layer\n", + "on the output to highlight how to swap in quantum modules during model construction." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7b5c2934", + "metadata": {}, + "outputs": [], + "source": [ + "# Install dependencies (run once)\n", + "!pip install -q jax jaxlib flax pennylane optax" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fd59f4ba", + "metadata": {}, + "outputs": [], + "source": [ + "import jax\n", + "import jax.numpy as jnp\n", + "from flax import linen as nn\n", + "import pennylane as qml\n", + "\n", + "print(\"JAX version\", jax.__version__)\n", + "print(\"Flax version\", nn.__version__)\n", + "print(\"PennyLane version\", qml.__version__)" + ] + }, + { + "cell_type": "markdown", + "id": "62a218ab", + "metadata": {}, + "source": [ + "## Define a simple quantum circuit and FlaxLayer" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "225690aa", + "metadata": {}, + "outputs": [], + "source": [ + "# one-qubit circuit that rotates by an input and a trainable weight\n", + "\n", + "dev = qml.device(\"default.qubit\", wires=1)\n", + "\n", + "@qml.qnode(dev, interface=\"jax\")\n", + "def qcircuit(inputs, weights):\n", + " # inputs expected shape (1,) and weights shape (1,)\n", + " qml.RX(inputs[0], wires=0)\n", + " qml.RY(weights[0], wires=0)\n", + " return qml.expval(qml.PauliZ(0))\n", + "\n", + "weight_shapes = {\"weights\": (1,)}\n", + "qlayer = qml.qnn.FlaxLayer(qcircuit, weight_shapes)\n", + "\n", + "# quick check\n", + "x = jnp.array([0.1])\n", + "w = jnp.array([0.2])\n", + "print(\"quantum output\", qlayer(x, w))" + ] + }, + { + "cell_type": "markdown", + "id": "fbaa57f9", + "metadata": {}, + "source": [ + "## Multi-Head Attention Module with Quantum Output" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5f8bc231", + "metadata": {}, + "outputs": [], + "source": [ + "class QuantumAttentionBlock(nn.Module):\n", + " dim: int\n", + " num_heads: int = 1\n", + "\n", + " def setup(self):\n", + " self.q_dense = nn.Dense(self.dim)\n", + " self.k_dense = nn.Dense(self.dim)\n", + " self.v_dense = nn.Dense(self.dim)\n", + " self.out_dense = nn.Dense(self.dim)\n", + " self.quantum = qlayer\n", + "\n", + " def __call__(self, x):\n", + " q = self.q_dense(x)\n", + " k = self.k_dense(x)\n", + " v = self.v_dense(x)\n", + " # simple scaled dot-product attention for demonstration\n", + " score = jnp.einsum(\"...d,...d->...\", q, k) / jnp.sqrt(self.dim)\n", + " weights = nn.softmax(score)\n", + " attended = weights[..., None] * v\n", + " classical = self.out_dense(attended)\n", + " # apply quantum layer to first feature vector in batch\n", + " qout = self.quantum(classical[..., 0], jnp.ones((1,)))\n", + " # broadcast quantum result and add to classical output\n", + " return classical + qout[..., None]\n", + "\n", + "# instantiate and run\n", + "model = QuantumAttentionBlock(dim=4)\n", + "params = model.init(jax.random.PRNGKey(0), jnp.ones((2,4)))\n", + "output = model.apply(params, jnp.ones((2,4)))\n", + "print(\"output shape\", output.shape)" + ] + }, + { + "cell_type": "markdown", + "id": "dd2013e2", + "metadata": {}, + "source": [ + "## Next Steps\n", + "\n", + "- Extend to genuine multi-head implementation.\n", + "- Swap `QuantumAttentionBlock` into larger transformer models.\n", + "- Benchmark against classical-only baselines using jet classification data.\n", + "\n", + "This notebook establishes the **basic pattern** for plugging a PennyLane `FlaxLayer`\n", + "into a JAX/Flax architecture. For QMLHEP7, similar bridges will be used inside a\n", + "particle transformer attention block." + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/Quantum_Attention_Benchmarks/README.md b/Quantum_Attention_Benchmarks/README.md new file mode 100644 index 0000000..74b969b --- /dev/null +++ b/Quantum_Attention_Benchmarks/README.md @@ -0,0 +1,19 @@ +# Quantum Attention Benchmarks + +This folder contains initial experiments for the **Quantum Particle Transformer** +(QMLHEP7) project. We focus on integrating quantum layers into attention-based +architectures using the JAX/Flax ecosystem. + +## Contents + +- `JAX_Quantum_Attention_Baseline.ipynb`: a minimal notebook demonstrating how to + replace classical dense layers in a simple attention block with a PennyLane + `FlaxLayer` quantum circuit. + +## Requirements + +- Python packages: `jax`, `jaxlib`, `flax`, `pennylane`, `optax`. +- A quantum device backend supported by PennyLane (e.g., `default.qubit`). + +This documentation improves visibility for researchers exploring Transformer-style +benchmarks within the QMLHEP repository. \ No newline at end of file