Quantum Machine Learning on Amazon Braket

Quantum Machine Learning on AWS

Quantum machine learning (QML) applies variational quantum circuits as trainable models. The circuit plays the role of a neural network layer: data is encoded into qubit rotations, parameterized gates are trained by gradient descent, and measurements produce predictions. Current NISQ devices cannot outperform classical ML on real tasks, but QML is an active research area and Braket provides a convenient cloud platform for experimentation.

Amazon Braket integrates natively with PennyLane through the amazon-braket-pennylane-plugin. This means you write standard PennyLane code and swap in a Braket device with one line: local simulation, managed cloud simulators (SV1, TN1, DM1), or real QPUs (IonQ, Rigetti, OQC, Oxford Quantum Circuits).

Setup

pip install amazon-braket-sdk amazon-braket-pennylane-plugin pennylane scikit-learn

Configure AWS credentials:

aws configure   # or set AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_DEFAULT_REGION

Local Simulation with PennyLane-Braket

Start with the local simulator before spending cloud credits:

import pennylane as qml
import numpy as np

# Local Braket simulator -- no AWS credentials or costs needed
dev_local = qml.device("braket.local.qubit", wires=4)

@qml.qnode(dev_local)
def circuit(x, weights):
    # Data encoding: angle encoding of 4 features
    for i in range(4):
        qml.RY(x[i], wires=i)

    # Variational layer 1
    for i in range(4):
        qml.RZ(weights[0, i], wires=i)
        qml.RY(weights[1, i], wires=i)

    # Entangling layer
    qml.CNOT(wires=[0, 1])
    qml.CNOT(wires=[1, 2])
    qml.CNOT(wires=[2, 3])
    qml.CNOT(wires=[3, 0])

    # Variational layer 2
    for i in range(4):
        qml.RZ(weights[2, i], wires=i)
        qml.RY(weights[3, i], wires=i)

    return qml.expval(qml.PauliZ(0))

# Test with random inputs
x_test = np.array([0.5, 1.2, 0.3, 0.8])
weights_test = np.random.uniform(-np.pi, np.pi, (4, 4))
print("Circuit output:", circuit(x_test, weights_test))

Loading and Preparing the Iris Dataset

The Iris dataset has 4 features and 3 classes. For binary classification we take the two most separable classes (setosa vs. versicolor, indices 0 and 1) and map labels to {+1, -1}.

from sklearn.datasets import load_iris
from sklearn.preprocessing import normalize, MinMaxScaler
from sklearn.model_selection import train_test_split
import numpy as np

iris = load_iris()
X = iris.data[:100]        # classes 0 and 1 only
y = iris.target[:100]
y = np.where(y == 0, 1.0, -1.0)  # map to {+1, -1}

# Scale features to [0, pi] for angle encoding
scaler = MinMaxScaler(feature_range=(0, np.pi))
X = scaler.fit_transform(X)

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

print(f"Training samples: {len(X_train)}, Test samples: {len(X_test)}")

Defining the Variational Classifier

import pennylane as qml
from pennylane import numpy as np

dev = qml.device("braket.local.qubit", wires=4)

@qml.qnode(dev, diff_method="parameter-shift")
def variational_classifier(x, weights, bias):
    # Angle encoding
    for i in range(4):
        qml.RY(x[i], wires=i)

    # Hardware-efficient ansatz -- 2 layers
    for layer in range(2):
        for i in range(4):
            qml.RZ(weights[layer, i, 0], wires=i)
            qml.RY(weights[layer, i, 1], wires=i)
            qml.RZ(weights[layer, i, 2], wires=i)
        qml.CNOT(wires=[0, 1])
        qml.CNOT(wires=[1, 2])
        qml.CNOT(wires=[2, 3])

    return qml.expval(qml.PauliZ(0))

def predict(X, weights, bias):
    return np.array([variational_classifier(x, weights, bias) + bias for x in X])

def cost(weights, bias, X, y):
    preds = predict(X, weights, bias)
    # Mean squared error loss
    return np.mean((preds - y) ** 2)

def accuracy(weights, bias, X, y):
    preds = np.sign(predict(X, weights, bias))
    return np.mean(preds == y)

Training with Gradient Descent

PennyLane computes gradients of variational_classifier via the parameter-shift rule, which is hardware-compatible (works identically on simulators and real QPUs).

from pennylane import numpy as np

np.random.seed(42)
weights = np.random.uniform(-np.pi, np.pi, (2, 4, 3), requires_grad=True)
bias = np.array(0.0, requires_grad=True)

opt = qml.AdamOptimizer(stepsize=0.05)

n_epochs = 30
batch_size = 10

print("Training variational classifier...")
for epoch in range(n_epochs):
    # Mini-batch stochastic gradient descent
    idx = np.random.choice(len(X_train), batch_size, replace=False)
    X_batch = X_train[idx]
    y_batch = y_train[idx]

    weights, bias, _, _ = opt.step(cost, weights, bias, X_batch, y_batch)

    if (epoch + 1) % 5 == 0:
        train_cost = cost(weights, bias, X_train, y_train)
        train_acc = accuracy(weights, bias, X_train, y_train)
        test_acc = accuracy(weights, bias, X_test, y_test)
        print(f"Epoch {epoch+1:3d} | Cost: {train_cost:.4f} | "
              f"Train acc: {train_acc:.2%} | Test acc: {test_acc:.2%}")

print(f"\nFinal test accuracy: {accuracy(weights, bias, X_test, y_test):.2%}")

On the local simulator, 30 epochs typically take 2-5 minutes and achieve 90-95% test accuracy on this binary Iris task.

Switching to the Braket SV1 Managed Simulator

SV1 is Braket’s managed state vector simulator. It scales to 34 qubits and runs faster than local simulation for larger circuits because it uses distributed compute.

import boto3

# Replace with your actual S3 bucket name
S3_BUCKET = "your-braket-bucket"
S3_PREFIX = "qml-tutorial"

dev_sv1 = qml.device(
    "braket.aws.qubit",
    device_arn="arn:aws:braket:::device/quantum-simulator/amazon/sv1",
    wires=4,
    s3_destination_folder=(S3_BUCKET, S3_PREFIX),
    shots=1000,  # SV1 with shots for realistic sampling
)

@qml.qnode(dev_sv1, diff_method="parameter-shift")
def classifier_sv1(x, weights, bias):
    for i in range(4):
        qml.RY(x[i], wires=i)
    for layer in range(2):
        for i in range(4):
            qml.RZ(weights[layer, i, 0], wires=i)
            qml.RY(weights[layer, i, 1], wires=i)
            qml.RZ(weights[layer, i, 2], wires=i)
        qml.CNOT(wires=[0, 1])
        qml.CNOT(wires=[1, 2])
        qml.CNOT(wires=[2, 3])
    return qml.expval(qml.PauliZ(0))

Each circuit evaluation on SV1 costs approximately $0.00075 per task plus$0.000145 per shot. For 30 epochs with 10 samples per batch and 3 gradient-shift evaluations per parameter (4x4x3 = 48 parameters, so 96 evaluations per step plus forward passes), expect roughly 30 * (10 * 97) = ~29,000 circuit evaluations total. At 1000 shots each that is approximately $4-6 for the full training run.

Hybrid Jobs for Persistent Training Loops

Amazon Braket Hybrid Jobs runs your classical-quantum training loop as a managed job, keeping a classical compute instance warm throughout training to avoid cold-start latency on each optimization step.

Create a job script train_job.py:

# train_job.py -- entry point for Braket Hybrid Job
import pennylane as qml
from pennylane import numpy as np
from braket.jobs import save_job_result
import json, os

hyperparams = json.loads(os.environ.get("AMZN_BRAKET_HP_FILE", "{}"))
n_epochs = int(hyperparams.get("n_epochs", 20))

dev = qml.device(
    "braket.aws.qubit",
    device_arn="arn:aws:braket:::device/quantum-simulator/amazon/sv1",
    wires=4,
    s3_destination_folder=(
        os.environ["AMZN_BRAKET_OUT_S3_BUCKET"],
        os.environ["AMZN_BRAKET_JOB_NAME"],
    ),
)

# ... (same circuit definition and training loop as above) ...

save_job_result({"final_accuracy": float(test_acc), "weights": weights.tolist()})

Submit from your local machine:

from braket.jobs import AwsQuantumJob

job = AwsQuantumJob.create(
    device="arn:aws:braket:::device/quantum-simulator/amazon/sv1",
    source_module="train_job.py",
    entry_point="train_job:main",
    hyperparameters={"n_epochs": "30"},
    instance_config={"instanceType": "ml.m5.large", "volumeSizeInGb": 30},
    job_name="iris-vqc-training",
)

print(f"Job ARN: {job.arn}")
print(f"Status: {job.state()}")

Hybrid Jobs automatically handle S3 result storage, CloudWatch logging, and IAM permissions. The warm classical instance eliminates the ~1 second cold-start overhead that accumulates painfully over thousands of circuit calls.

Summary and Costs

Simulator	Max Qubits	Speed (4q, 1000 shots)	Cost
braket.local.qubit	~25	~0.3 s/circuit	Free
SV1 (managed)	34	~0.8 s/circuit	~$0.001/circuit
TN1 (tensor network)	50	Circuit-dependent	~$0.001/circuit

For initial development and debugging, always use the local simulator. Switch to SV1 when you need shot-based noise, multi-shot statistics, or want reproducible cloud results for a paper or production system. Use Hybrid Jobs when your training loop runs for more than a few minutes to avoid repeated cold-start overhead.

Next Steps

• Try qml.device("braket.aws.qubit", device_arn="arn:aws:braket:us-east-1::device/qpu/ionq/Aria-1") to run your trained circuit on real IonQ hardware.
• Explore qml.kernels for quantum kernel methods, which bypass gradient computation entirely by measuring inner products between quantum feature states.
• Use the Braket Cost Tracker in the AWS console to monitor spending during experiments.