Multidimensional regression with a variational quantum circuit¶

What is a quantum model?¶

A quantum model $g_{\vec{\theta}}(\vec{x})$ is the expectation value of some observable $M$ estimated on the state prepared by a parameterized circuit $U(\vec{x}, \vec{\theta})$:

By repeatedly running the circuit with a set of parameters $\vec{\theta}$ and set of data points $\vec{x}$, we can approximate the expectation value of the observable $M$ in the state $U(\vec{x}, \vec{\theta}) | 0 \rangle.$ Then, the parameters can be optimized to minimize some loss function.

Building the variational circuit¶

In this example, we will use a variational quantum circuit to find the Fourier series that approximates the function $f(x_1, x_2) = \frac{1}{2} \left( x_1^2 + x_2^2 \right)$. The variational circuit that we are using is made up of $L$ layers. Each layer consists of a data-encoding block $S(\vec{x})$ and a training block $W(\vec{\theta})$. The overall circuit is:

The training blocks $W(\vec{\theta})$ depend on the parameters $\vec{\theta}$ that can be optimized classically.

We will build a circuit such that the expectation value of the $Z\otimes Z$ observable is a partial Fourier series that approximates $f(\vec{x})$, i.e.,

Then, we can directly plot the partial Fourier series. We can also apply a Fourier transform to $g_{\vec{\theta}}$, so we can obtain the Fourier coefficients, $c_\vec{\omega}$. To know more about how to obtain the Fourier series, check out these two related tutorials 2, 3.

Constructing the quantum circuit¶

First, let’s import the necessary libraries and seed the random number generator. We will use Matplotlib for plotting and JAX 4 for optimization. We will also define the device, which has two qubits, using device().

import matplotlib.pyplot as plt
import pennylane as qml
from pennylane import numpy as pnp
import jax
from jax import numpy as jnp
import optax

pnp.random.seed(42)

dev = qml.device('default.qubit', wires=2)

Now we will construct the data-encoding circuit block, $S(\vec{x})$, as a product of $R_z$ rotations:

Specifically, we define the $S(\vec{x})$ operator using the AngleEmbedding function.

def S(x):
    qml.AngleEmbedding( x, wires=[0,1],rotation='Z')

For the $W(\vec{\theta})$ operator, we will use an ansatz that is available in PennyLane, called StronglyEntanglingLayers.

def W(params):
    qml.StronglyEntanglingLayers(params, wires=[0,1])

Now we will build the circuit in PennyLane by alternating layers of $W(\vec{\theta})$ and $S(\vec{x})$ layers. On this prepared state, we estimate the expectation value of the $Z\otimes Z$ operator, using PennyLane’s expval() function.

@qml.qnode(dev,interface="jax")
def quantum_neural_network(params, x):
    layers=len(params[:,0,0])-1
    n_wires=len(params[0,:,0])
    n_params_rot=len(params[0,0,:])
    for i in range(layers):
      W(params[i,:,:].reshape(1,n_wires,n_params_rot))
      S(x)
    W(params[-1,:,:].reshape(1,n_wires,n_params_rot))

    return qml.expval(qml.PauliZ(wires=0)@qml.PauliZ(wires=1))

The function we will be fitting is $f(x_1, x_2) = \frac{1}{2} \left( x_1^2 + x_2^2 \right)$, which we will define as target_function:

def target_function(x):
    f=1/2*(x[0]**2+x[1]**2)
    return f

Now we will specify the range of $x_1$ and $x_2$ values and store those values in an input data vector. We are fitting the function for $x_1, x_2 \in [-1, 1]$ using 30 evenly spaced samples for each variable.

x1_min=-1
x1_max=1
x2_min=-1
x2_max=1
num_samples=30

Now we build the training data with the exact target function $f(x_1, x_2)$. To do so, it is convenient to create a two-dimensional grid to make sure that, for each value of $x_1,$ we perform a sweep over all the values of $x_2$ and viceversa.

x1_train=pnp.linspace(x1_min,x1_max, num_samples)
x2_train=pnp.linspace(x2_min,x2_max, num_samples)
x1_mesh,x2_mesh=pnp.meshgrid(x1_train, x2_train)

We define x_train and y_train using the above vectors, reshaping them for our convenience

x_train=pnp.stack((x1_mesh.flatten(), x2_mesh.flatten()), axis=1)
y_train = target_function([x1_mesh,x2_mesh]).reshape(-1,1)
# Let's take a look at how they look like
print("x_train:\n", x_train[:5])
print("y_train:\n", y_train[:5])

x_train:
 [[-1.         -1.        ]
 [-0.93103448 -1.        ]
 [-0.86206897 -1.        ]
 [-0.79310345 -1.        ]
 [-0.72413793 -1.        ]]
y_train:
 [[1.        ]
 [0.9334126 ]
 [0.87158145]
 [0.81450654]
 [0.76218787]]

Optimizing the circuit¶

We want to optimize the circuit above so that the expectation value of $Z \otimes Z$ approximates the exact target function. This is done by minimizing the mean squared error between the circuit output and the exact target function. In particular, the optimization process to train the variational circuit will be performed using JAX, an auto differentiable machine learning framework to accelerate the classical optimization of the parameters. Check out 4 to learn more about how to use JAX to optimize your QML models.

@jax.jit
def mse(params,x,targets):
    # We compute the mean square error between the target function and the quantum circuit to quantify the quality of our estimator
    return (quantum_neural_network(params,x)-jnp.array(targets))**2
@jax.jit
def loss_fn(params, x,targets):
    # We define the loss function to feed our optimizer
    mse_pred = jax.vmap(mse,in_axes=(None, 0,0))(params,x,targets)
    loss = jnp.mean(mse_pred)
    return loss

Here, we are choosing an Adam optimizer with a learning rate of 0.05 and 300 steps.

opt = optax.adam(learning_rate=0.05)
max_steps=300

@jax.jit
def update_step_jit(i, args):
    # We loop over this function to optimize the trainable parameters
    params, opt_state, data, targets, print_training = args
    loss_val, grads = jax.value_and_grad(loss_fn)(params, data, targets)
    updates, opt_state = opt.update(grads, opt_state)
    params = optax.apply_updates(params, updates)

    def print_fn():
        jax.debug.print("Step: {i}  Loss: {loss_val}", i=i, loss_val=loss_val)
    # if print_training=True, print the loss every 50 steps
    jax.lax.cond((jnp.mod(i, 50) == 0 ) & print_training, print_fn, lambda: None)
    return (params, opt_state, data, targets, print_training)

@jax.jit
def optimization_jit(params, data, targets, print_training=False):
    opt_state = opt.init(params)
    args = (params, opt_state, jnp.asarray(data), targets, print_training)
    # We loop over update_step_jit max_steps iterations to optimize the parameters
    (params, opt_state, _, _, _) = jax.lax.fori_loop(0, max_steps+1, update_step_jit, args)
    return params

Now we will train the variational circuit with 4 layers and obtain a vector $\vec{\theta}$ with the optimized parameters.

wires=2
layers=4
params_shape = qml.StronglyEntanglingLayers.shape(n_layers=layers+1,n_wires=wires)
params=pnp.random.default_rng().random(size=params_shape)
best_params=optimization_jit(params, x_train, jnp.array(y_train), print_training=True)

Step: 0  Loss: 0.2148337960243225
Step: 50  Loss: 0.0238615944981575
Step: 100  Loss: 0.004363568499684334
Step: 150  Loss: 0.0029268567450344563
Step: 200  Loss: 0.002278682542964816
Step: 250  Loss: 0.0020085107535123825
Step: 300  Loss: 0.0016955287428572774

If you run this yourself, you’ll see that the training step with JAX is extremely fast! Once the optimized $\vec{\theta}$ has been obtained, we can use those parameters to build our fitted version of the function.

def evaluate(params, data):
    y_pred = jax.vmap(quantum_neural_network, in_axes=(None, 0))(params, data)
    return y_pred
y_predictions=evaluate(best_params,x_train)

To compare the fitted function to the exact target function, let’s take a look at the $R^2$ score:

from sklearn.metrics import r2_score
r2 = round(float(r2_score(y_train, y_predictions)),3)
print("R^2 Score:", r2)

R^2 Score: 0.967

Let’s now plot the results to visually check how good our fit is!

fig = plt.figure()
# Target function
ax1 = fig.add_subplot(1, 2, 1, projection='3d')
ax1.plot_surface(x1_mesh,x2_mesh, y_train.reshape(x1_mesh.shape), cmap='viridis')
ax1.set_zlim(0,1)
ax1.set_xlabel('$x$',fontsize=10)
ax1.set_ylabel('$y$',fontsize=10)
ax1.set_zlabel('$f(x,y)$',fontsize=10)
ax1.set_title('Target ')

# Predictions
ax2 = fig.add_subplot(1, 2, 2, projection='3d')
ax2.plot_surface(x1_mesh,x2_mesh, y_predictions.reshape(x1_mesh.shape), cmap='viridis')
ax2.set_zlim(0,1)
ax2.set_xlabel('$x$',fontsize=10)
ax2.set_ylabel('$y$',fontsize=10)
ax2.set_zlabel('$f(x,y)$',fontsize=10)
ax2.set_title(f' Predicted \nAccuracy: {round(r2*100,3)}%')

# Show the plot
plt.tight_layout(pad=3.7)

Cool! We have managed to successfully fit a multidimensional function using a parametrized quantum circuit!

Conclusions¶

In this demo, we’ve shown how to utilize a variational quantum circuit to solve a regression problem for a two-dimensional function. The results show a good agreement with the target function and the model can be trained further, increasing the number of iterations in the training to maximize the accuracy. It also paves the way for addressing a regression problem for an $N$-dimensional function, as everything presented here can be easily generalized. A final check that could be done is to obtain the Fourier coefficients of the trained circuit and compare it with the Fourier series we obtained directly from the target function.

References¶

1

Maria Schuld, Ryan Sweke, and Johannes Jakob Meyer “The effect of data encoding on the expressive power of variational quantum machine learning models.”, arXiv:2008.0865, 2021.

2

Johannes Jakob Meyer, Maria Schuld “Tutorial: Quantum models as Fourier series”, Pennylane: Quantum models as Fourier series, 2021.

3

Jorge J. Martinez de Lejarza “Tutorial: Quantum Fourier Iterative Amplitude Estimation”, Qibo: Quantum Fourier Iterative Amplitude Estimation, 2023.

4(1,2)

Josh Izaac, Maria Schuld “How to optimize a QML model using JAX and Optax”, Pennylane: How to optimize a QML model using JAX and Optax, 2024

About the authors¶

Total running time of the script: (0 minutes 4.116 seconds)