Quantum natural gradient optimizer with user-defined gates

Hi,

I was trying to use the quantum natural gradient (QNG) optimizer on a toy VQE problem where the ansatz employs user-defined gates ‘XX’ and ‘YY’ (defined below).

When trying to take an optimizer step, I got the following error: “ValueError: solve1: Input operand 1 has a mismatch in its core dimension 0, with gufunc signature (m,m),(m)->(m) (size 3 is different from 4)” coming from

np.linalg.solve(self.metric_tensor, grad_flat)


I didn’t encounter this error when using a circuit with Pennylane-defined gates. I think the error might be because I use the same parameter for both XX and YY gates but wasn’t sure how to fix this error.

This was my code:

from pennylane import numpy as np
import pennylane as qml

dev = qml.device("default.qubit", wires=2)
dev.operations.update({"XX", "YY"})

def circuit(params, wires=[0, 1]):
qml.PauliX(wires=wires[1])
qml.RZ(params[0], wires=wires[0])
qml.RZ(params[1], wires=wires[1])
XX(params[2], wires=wires)
YY(params[2], wires=wires)

coeffs = [1, 1]
obs = [qml.PauliX(0), qml.PauliZ(1)]
H = qml.Hamiltonian(coeffs, obs)
cost_fn = qml.ExpvalCost(circuit, H, dev)

init_params = np.random.uniform(size=3)

opt = qml.QNGOptimizer(stepsize=0.01, diag_approx=False)

params, prev_energy = opt.step_and_cost(cost_fn, init_params)


And the two user-defined gates are:

class XX(qml.operation.Operation):
num_params = 1
num_wires = 2
par_domain = "R"
generator = [np.kron(X, X), -0.5]

@staticmethod
def _matrix(*params):
return np.cos(params[0]/2)*np.eye(4) + 1.j*np.sin(params[0]/2)*np.kron(X,X)

class YY(qml.operation.Operation):
num_params = 1
num_wires = 2
par_domain = "R"
generator = [np.kron(Y, Y), -0.5]

@staticmethod
def _matrix(*params):
return np.cos(params[0]/2)*np.eye(4) + 1.j*np.sin(params[0]/2)*np.kron(Y,Y)

Hi @hsim!

The issue here is that the metric tensor is only defined for the quantum circuit gate arguments, while the optimizer is instead optimizing the QNode arguments.

This is a bit of subtlety—in most cases, there is a 1-1 mapping between QNode arguments and gate arguments, so this issue doesn’t arise! As you note, however, the repeated parameter between the XX and YY breaks this 1-1 mapping.

In other words: the optimizer is assuming that the object it is optimizing is a pure quantum circuit with a well-defined metric tensor with respect to QNode arguments, which is not the case.

Workarounds

One workaround is to write out the QNG optimization manually using qml.metric_tensor (which extracts the metric tensor of the underlying quantum circuit), but extend it to take into account classical processing between QNode args and gate args.

This is pretty simple to do.

Let

• x\in \mathbb{R}^m be the QNode parameters

• f: \mathbb{R}^m \rightarrow \mathbb{R}^n be the function representing the transformation from QNode arguments to gate arguments,

• q: \mathbb{R}^n \rightarrow \mathbb{R} be the quantum function,

• C = q \circ f: \mathbb{R}^m\rightarrow \mathbb{R} the overall cost function, and

• g^{-1} be the pseudo-inverse of the Fubini-Study metric tensor of q.

Our cost function is therefore C(x) = q(f(x)), with a (quantum) natural gradient of

We can use the qml.transforms.classical_jacobian() transform to extract \mathbf{J}_f(x), qml.metric_tensor() to extract g, and the (not well documented) qnode.qtape.jacobian method to extract \mathbf{J}_q(x):

from pennylane import numpy as np
import pennylane as qml
from scipy.linalg import pinvh

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

@qml.qnode(dev)
def circuit(params, wires=[0, 1]):
qml.RZ(params[0], wires=wires[0])
qml.RZ(params[0], wires=wires[1])
qml.CNOT(wires=[0, 1])
qml.RZ(params[1], wires=wires[1])
qml.CRX(params[2], wires=wires)
qml.CRY(params[2], wires=wires)
return qml.expval(qml.PauliY(0))

init_params = np.array([0.1, 0.2, 0.3])

# extract Jf(x)
Jf = qml.transforms.classical_jacobian(circuit)(init_params)

# extract Jq(x)
circuit.construct([init_params], {})
Jq = circuit.qtape.jacobian(params=Jf @ init_params, device=dev)

# extract g
g= qml.metric_tensor(circuit)(init_params)

print("NatGad =", Jq @ pinvh(g) @ Jf)


Or, writing this as a function:

def natgrad(qnode):
Jf = qml.transforms.classical_jacobian(qnode)
g= qml.metric_tensor(circuit)
dev = qnode.device

qnode.construct([params], {})
gate_params = Jf(params) @ params
Jq = circuit.qtape.jacobian(params=gate_params, device=dev)
return Jq @ pinvh(g(params)) @ Jf(params)



Note: in the case where J_f = I, this reduces down to our standard natural gradient result:

@qml.qnode(dev)
def circuit(params, wires=[0, 1]):
qml.RZ(params[0], wires=wires[0])
# qml.RZ(params[0], wires=wires[1])
qml.CNOT(wires=[0, 1])
qml.RZ(params[1], wires=wires[1])
qml.CRX(params[2], wires=wires)
# qml.CRY(params[2], wires=wires)
return qml.expval(qml.PauliY(0))

g = qml.metric_tensor(circuit)(init_params)


1. We should add a qml.transforms.quantum_jacobian function to make it easier to extract the quantum portion of the Jacobian directly

2. Long term, it may make sense to build this directly into PennyLane. That is, re-position QNG as a gradient method, rather than an optimization technique.

1 Like

Thanks, @josh!

How would I extend this for qml.ExpvalCost which, if I understood correctly, uses a collection of QNodes?

No worries @hsim!

I’ve generalized my previous answer to account for ExpvalCost:

def quantum_jacobian(qnode):
"""Returns a function that computes the purely
quantum Jacobian of a QNode"""
if isinstance(qnode, qml.ExpvalCost):
qnodes = qnode.qnodes
coeffs = qnode.hamiltonian.coeffs

def _jacobian(*args, **kwargs):
jacs = [quantum_jacobian(q)(*args, **kwargs) for q in qnodes]
return qml.math.sum([c * j for c, j in zip(coeffs, jacs)], axis=0)

return _jacobian

dev = qnode.device

def _jacobian(*args, **kwargs):
qnode.construct(args, kwargs)
params = qml.math.stack(qnode.qtape.get_parameters())

def func(params):
return qml.math.stack(qnode.qtape.execute(params=params, device=dev))

return qml.jacobian(func)(params)

return _jacobian

"""Returns a function that computes
the natural gradient of a QNode"""
Jq = quantum_jacobian(qnode)
g = qml.metric_tensor(qnode)

if isinstance(qnode, qml.ExpvalCost):
qnode = qnode.qnodes[0]

Jf = qml.transforms.classical_jacobian(qnode)

mt = g(*args, **kwargs)
c = Jf(*args, **kwargs)
q = Jq(*args, **kwargs)
ng = qml.math.dot(q, pinvh(mt))
return qml.math.squeeze(qml.math.dot(ng, c))



With these functions defined, the following now works:

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

def circuit(params, wires=[0, 1]):
qml.RZ(params[0], wires=wires[0])
qml.RZ(params[0], wires=wires[1])
qml.CNOT(wires=[0, 1])
qml.RZ(params[1], wires=wires[1])
qml.CRX(params[2], wires=wires)
qml.CRY(params[2], wires=wires)

coeffs = [1, 1]
obs = [qml.PauliX(0), qml.PauliZ(1)]
H = qml.Hamiltonian(coeffs, obs)

circuit = qml.ExpvalCost(circuit, H, dev, interface="autograd")
init_params = np.array([0.1, 0.2, 0.3])

>>> print("Cost function=", circuit(init_params))
Cost function= 0.8501644188535566