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

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
for k in range(n_features):
qml.CSWAP(wires=[0, k + 1, 2 + k + 1])

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.preprocessing import StandardScaler
from sklearn.decomposition import PCA
# Load and preprocess the dataset
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)

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]

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"Step {i+1} cost = {cost_val}")
# reshape
pars = new_flat_pars.reshape(init_pars.shape)
$$$$


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
for k in range(n_features):
qml.CSWAP(wires=[0, k + 1, 2 + k + 1])

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
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.preprocessing import StandardScaler
from sklearn.decomposition import PCA
# Load and preprocess the dataset
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)

opt_state = optimizer.init(theta)

energy = []
thetas = []

@jax.jit

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)
opt_state, theta = partial_step(grad_circuit, opt_state, theta)

energy.append(val)
`