Quantum Optimization example for Tensorflow with optimizer.

In my last attempt to do a simple qubit rotation with Tensorflow-Quantum (TFQ) I struggled to implement the automatic gradient update with TFQ’s optimizers. I show now a nice way of how we can use the gradient descent optimizer of TF. I will demonstrate this on a variational circuit V(θ) to rotate the initial state |0000> into the equal superposition state ∑|σ>, where σ are all the possible qubit configurations. The code can be found on GitHub.

One layer of the variational circuit contains a general rotation 𝑅(𝜙,𝜃,𝜔)=𝑅𝑧(𝜔)𝑅𝑧(𝜃)𝑅𝑧(𝜙) on each qubit, followed by entangling gates. We apply 2 layers and use 4 qubits.

Variational Circuit V

I will present two methods of how we can do this. The first one was pushed to my Github repository by Alejandro Pozas. He uses the keras functions to compile and fit the circuit. This version is very compact. The second version is my own one updated in a way that we can use the TF optimizers. It is more basic TF and if you are also more familiar with pytorch you probably will understand this version better.

First we initialize the circuit with Cirq:

def generate_circuit(nr_of_qubits, layers):
    qubits = cirq.GridQubit.rect(1, nr_of_qubits) 
    nr_parameters = 3*nr_of_qubits*layers 
   
    symb = sympy.symbols('theta0:'+str(nr_parameters))
    symbols = np.array(symb)
    symbols = symbols.reshape(layers, nr_of_qubits, 3)
    circuit = cirq.Circuit()
    for l in range(layers):
        for i, qubit in enumerate(qubits):
            circuit += cirq.rz(symbols[l][i][0])(qubit)
            circuit += cirq.rx(symbols[l][i][1])(qubit)
            circuit += cirq.rz(symbols[l][i][2])(qubit)        
circuit += cirq.CZ(qubits[0], qubits[1])
        circuit += cirq.CZ(qubits[2], qubits[3])
        circuit += cirq.CZ(qubits[1], qubits[2]     
    
    op = 1/4*(cirq.X(qubits[0]) + cirq.X(qubits[1]) + 
              cirq.X(qubits[2]) + cirq.X(qubits[3]))         
    
    return circuit, op, list(symb)
nr_of_qubits = 4
layers = 2
tf_circuit, op, (symbols) = generate_circuit(nr_of_qubits, layers)
SVGCircuit(tf_circuit)

This function also returns the operator 𝑀=1/4∗(𝑋1+𝑋2+𝑋3+𝑋4), which measures the fidelity with the target state. In the end we want to minimize
1-<M>.

Method without Keras (pure TF)

For the pure TF method we define a model that will return the quantum circuit’s parameters. In the code snippet here I show the most basic model that only forwards its parameters. But this could be any possible function.

class circuit_params(Model):
    def __init__(self, W = np.random.normal(2)):
        super(circuit_params, self).__init__()
        self.W = tf.Variable(W)
 
    @tf.function
    def __call__(self):
        return self.W
model = circuit_params(W = np.random.rand(nr_of_qubits*layers*3))

If we call model() it will simply return the TF variable W which we will in a next step use as the circuit parameters. The training routine looks the following.

@tf.function
def train_step():
    with tf.GradientTape() as tape:
        params = model()
        out = tfq.layers.Expectation()(tf_circuit,
                                       symbol_names=symbols,
                                       symbol_values=[params],
                                       operators=1-op)
        
    gradients = tape.gradient(out, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients,       
                                  model.trainable_variables))

As stated before we get the circuit parameters params from the classical model object and then we use the same parameters to calculate the expectation value of our circuit with respect to the operator op that we defined before.

Finally we can iterate through the training steps several times until the output converges.

for i in range(100):
    train_step()

Method with Keras

To do the same with TFQ and Keras we first convert the Cirq circuit into a TFQ layer and define the measurement operator which we would like to optimize.

outputs = tfq.layers.PQC(tf_circuit,         
                         1-op)

Then we define a sequential model with an input of arbitrary dimension and the measurement from the quantum layer as an output.

model = tf.keras.Sequential([
    tf.keras.layers.Input(shape=(), dtype=tf.string),
    outputs
])

To compile the model we need to be aware that the keras needs a loss function of the form f(labels, predictions) for sequential models, but we don’t have labels. To fix this problem we can use a loss function with two arguments for the labels and the predictions that simply returns the prediction, which in our case is the measurement 1-<M> that we want to minimize.

def loss(real, pred):
    return pred

After defining the loss we can compile the model.

model.compile(loss=loss,
              optimizer=tf.keras.optimizers.SGD(learning_rate=0.4))

Before we can fit the model we define the input to our model, which is the Cirq circuit converted to a TF tensor. And we also define an empty vector for our dummy labels.

dummy_input  = tfq.convert_to_tensor([cirq.Circuit()])
dummy_output = np.array([[]])

And finally we can fit the model.

steps = 100
model.fit(dummy_input, dummy_output, epochs=steps)