Optimizing a parametrized Quantum Circuit in batches does not decrease the cost function while unbatched optimization does

Question

I want to optimize a variational quantum circuit to maximize the Hilbert-Schmidt Distance between the different classes of the UCI breast cancer data set. When I choose to use batched optimization, the circuit does not really get optimized. That means the values cost function does not really decrease and always stay randomly on an interval of approximately [0.8, 0.9].

I don't think that the problem is a barren plateau, because the gradients are not too small and when I instead use a small sample of the data and try to run all optimization steps on the same data, instead of randomly sampled batches, the evaluations of the cost function are decreasing in a reasonable manner.

Does anybody here see, what is the problem of the batched optimization?

All my code available here: https://github.com/Rlag1998/Embedding_Generalization/blob/main/tutorial_embedding_generalization.ipynb

I made some changes to adjust the code to the most recent version of Pennylane. I also deleted the classical linear layer before the circuit, as i was thinking that most of the training took place in there instead inside the quantum circuit.

This is my code:

import pennylane as qml
from pennylane import numpy as np
from pennylane import RX, RY, RZ, CNOT

np.random.seed(seed=1234)

n_features = 2
n_qubits = 2 * n_features + 1

dev = qml.device("default.qubit", wires=n_qubits)


def feature_encoding_hamiltonian(features, wires):

    for idx, w in enumerate(wires):
        RX(features[idx], wires=w)


def ising_hamiltonian(weights, wires, l):

    # ZZ coupling
    CNOT(wires=[wires[1], wires[0]])
    RZ(weights[l, 0], wires=wires[0])
    CNOT(wires=[wires[1], wires[0]])
    # local fields
    for idx, w in enumerate(wires):
        RY(weights[l, idx + 1], wires=w)


def QAOAEmbedding(features, weights, wires):

    repeat = len(weights)
    for l in range(repeat):
        # apply alternating Hamiltonians
        feature_encoding_hamiltonian(features, wires)
        ising_hamiltonian(weights, wires, l)
    # repeat the feature encoding once more at the end
    feature_encoding_hamiltonian(features, wires)

@qml.qnode(dev, argnum=0)
def swap_test(q_weights, x1, x2):

    # load the two inputs into two different registers
    QAOAEmbedding(features=x1, weights=q_weights, wires=[1, 2])
    QAOAEmbedding(features=x2, weights=q_weights, wires=[3, 4])

    # perform the SWAP test
    qml.Hadamard(wires=0)
    for k in range(n_features):
        qml.CSWAP(wires=[0, k + 1, 2 + k + 1])
    qml.Hadamard(wires=0)

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


def overlaps(weights, X1=None, X2=None):
    overlap = 0
    for x1 in X1:
        for x2 in X2:
            # overlap of embedded intermediate features
            overlap += swap_test(q_weights=weights, x1=x1, x2=x2)

    mean_overlap = overlap / (len(X1) * len(X2))

    return mean_overlap

def cost(weights, A=None, B=None):
    aa = overlaps(weights, X1=A, X2=A)
    bb = overlaps(weights, X1=B, X2=B)
    ab = overlaps(weights, X1=A, X2=B)

    d_hs = -2 * ab + (aa + bb)
    #print("print cost in cost func: ", 1-0.5 * d_hs)
    return 1 - 0.5 * d_hs

from sklearn.datasets import load_breast_cancer
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
# Load and preprocess the dataset
data = load_breast_cancer()
X = data.data
y = data.target

# Standardize the features
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# Reduce dimensions with PCA
pca = PCA(n_components=2)
X_pca = pca.fit_transform(X_scaled)

# Split the dataset into two classes
A = X_pca[y == 0]
B = X_pca[y == 1]

# size=(x, y) x defines the number of layers of the circuit and y the no of 
# trainable parameters in each layer
init_pars = np.random.normal(loc=0, scale=0.1, size=(1, 3), requires_grad=True)

optimizer = qml.AdamOptimizer(stepsize=0.05)
batch_size = 10
pars = init_pars

for i in range(100):
    # Sample a batch of training inputs from each class
    selectA = np.random.choice(range(len(A)), size=(batch_size,), replace=True)
    selectB = np.random.choice(range(len(B)), size=(batch_size,), replace=True)
    A_batch = [A[s] for s in selectA]
    B_batch = [B[s] for s in selectB]

    #grad = qml.grad(cost)(pars, A_batch, B_batch)
    flat_pars = pars.flatten()
    new_flat_pars, cost_val = optimizer.step_and_cost(lambda w: cost(w.reshape(init_pars.shape), A=A_batch, B=B_batch), flat_pars)
    #print(f"cost = {cost_val}, gradient norm = {np.linalg.norm(grad):.4f}")
    print(f"Step {i+1} cost = {cost_val}")
    # reshape
    pars = new_flat_pars.reshape(init_pars.shape)
```

Answer 1

It seems the batch size is too small (what constitutes too small batch size is a science of its own). You need to use larger batch sizes but that comes at a performance cost. Luckily there is a way out :)

Your cost function involves a two-fold for loop over the two batches and that takes a really long time in python. Instead, you should vectorize your code.

For me the easiest to do this is with jax.vmap, which has the additional benefit that I can just-in-time compile the function (so every consecutive call in the training is using compiled code).

I took the liberty of modifying your example to run vectorized with jax, its much faster now and helps with trying different epochs and batch sizes, though unfortunately I could not dig deeper into playing with the model.

Hope this helps!

import pennylane as qml
from pennylane import RX, RY, RZ, CNOT
import jax
import jax.numpy as jnp
import optax
import numpy as np

n_features = 2
n_qubits = 2 * n_features + 1

dev = qml.device("default.qubit", wires=n_qubits)


def feature_encoding_hamiltonian(features, wires):

    for idx, w in enumerate(wires):
        RX(features[idx], wires=w)


def ising_hamiltonian(weights, wires, l):

    # ZZ coupling
    CNOT(wires=[wires[1], wires[0]])
    RZ(weights[l, 0], wires=wires[0])
    CNOT(wires=[wires[1], wires[0]])
    # local fields
    for idx, w in enumerate(wires):
        RY(weights[l, idx + 1], wires=w)


def QAOAEmbedding(features, weights, wires):

    repeat = len(weights)
    for l in range(repeat):
        # apply alternating Hamiltonians
        feature_encoding_hamiltonian(features, wires)
        ising_hamiltonian(weights, wires, l)
    # repeat the feature encoding once more at the end
    feature_encoding_hamiltonian(features, wires)

@qml.qnode(dev, interface="jax")
def swap_test(q_weights, x1, x2):

    # load the two inputs into two different registers
    QAOAEmbedding(features=x1, weights=q_weights, wires=[1, 2])
    QAOAEmbedding(features=x2, weights=q_weights, wires=[3, 4])

    # perform the SWAP test
    qml.Hadamard(wires=0)
    for k in range(n_features):
        qml.CSWAP(wires=[0, k + 1, 2 + k + 1])
    qml.Hadamard(wires=0)

    return qml.expval(qml.PauliZ(0))
def overlaps(weights, X1=None, X2=None):
    vector_swap = jax.vmap(jax.vmap(swap_test, (None, 0, None)), (None, None, 0))

    return jnp.mean(vector_swap(weights, X1, X2))

@jax.jit
@jax.value_and_grad
def value_and_grad(weights, A=None, B=None):
    aa = overlaps(weights, X1=A, X2=A)
    bb = overlaps(weights, X1=B, X2=B)
    ab = overlaps(weights, X1=A, X2=B)

    d_hs = -2 * ab + (aa + bb)
    #print("print cost in cost func: ", 1-0.5 * d_hs)
    return 1 - 0.5 * d_hs

from sklearn.datasets import load_breast_cancer
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
# Load and preprocess the dataset
data = load_breast_cancer()
X = data.data
y = data.target

# Standardize the features
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# Reduce dimensions with PCA
pca = PCA(n_components=2)
X_pca = pca.fit_transform(X_scaled)

# Split the dataset into two classes
A = X_pca[y == 0]
B = X_pca[y == 1]


theta = jax.random.normal(jax.random.PRNGKey(0), shape=(1, 3))
n_epochs = 10
batch_size = 1000
keys = jax.random.split(jax.random.PRNGKey(0), n_epochs)

optimizer = optax.adam(learning_rate=0.05)
opt_state = optimizer.init(theta)

energy = []
gradients = []
thetas = []

@jax.jit
def partial_step(grad_circuit, opt_state, theta):
    updates, opt_state = optimizer.update(grad_circuit, opt_state)
    theta = optax.apply_updates(theta, updates)

    return opt_state, theta


## Optimization loop
for n in range(n_epochs):
    selectA = np.random.choice(range(len(A)), size=(batch_size,), replace=True)
    selectB = np.random.choice(range(len(B)), size=(batch_size,), replace=True)
    A_batch = jnp.array([A[s] for s in selectA])
    B_batch = jnp.array([B[s] for s in selectB])

    # val, theta, grad_circuit, opt_state = step(theta, opt_state)
    val, grad_circuit = value_and_grad(theta, A_batch, B_batch)
    opt_state, theta = partial_step(grad_circuit, opt_state, theta)

    energy.append(val)
    gradients.append(
        grad_circuit
    )
    thetas.append(
        theta
    )
import matplotlib.pyplot as plt
plt.plot(energy)

Here is an example with batch size 1000 trainer over 10 epochs: