Intro to Amplitude Amplification

The algorithm

To obtain the state \(|\phi\rangle,\) we could just rotate the initial state counterclockwise by an angle \(\pi/2 -\theta.\) However, we don’t explicitly know \(|\phi\rangle,\) so it’s unclear how this could be done. This is where a great idea is born: what if instead of rotations we think of reflections?

The main insight of the Amp Amp algorithm is that there is a sequence of two reflections that allow us to effectively perform a rotation towards the target state. The first is the reflection with respect to \(|\phi^{\perp}\rangle\) and the second one is the reflection with respect to \(|\Psi\rangle.\)

Let’s go step by step. First we apply the reflection around \(|\phi^{\perp}\rangle:\)

This reflection may seem challenging to implement since we do not explicitly know \(|\phi^{\perp}\rangle.\) However, the operator performing the reflection is well-defined:

Amp Amp usually assumes access to an oracle implementing this reflection. For example, in a search problem, this is the usual Grover oracle that “marks” the target state with a phase flip. In practice, the way to build the oracle is just a classic check: if the given state meets the known property, we change its sign. This will become clearer later with an example.

After applying the first reflection we are moving away from \(|\phi\rangle\) — why do that? Well, sometimes it’s necessary to take a step backward to take two steps forward, and that is exactly what we will do. For this purpose, we use a second reflection with respect to \(|\Psi\rangle.\) This is easier to build since we know the operator \(U\) that generates \(|\Psi\rangle.\)

Together, these two reflections are equivalent to rotating the state by \(2\theta\) degrees from its original position, where \(\theta\) is the angle that defines the initial state. To amplify the amplitude and approach the target state, we perform this sequence of rotations multiple times, i.e. \(\dots R_{\Psi}R_{\phi^{\perp}}R_{\Psi}R_{\phi^{\perp}}.\) More precisely, we repeat them \(k\) times, with \(k\) given by:

This expression is derived by recognizing that the angle of the resulting state after \(k\) iterations is \((2k + 1)\theta,\) and we aim for this value to be equal to \(\frac{\pi}{2}\) radians (i.e. \(90º\)).

As we will see below, Amp Amp can be applied to unstructured dataset searching problems. Let’s suppose that in a set of N elements we are looking for a single one, and we begin with an equal superposition state such that \(\alpha=\frac{1}{\sqrt{N}}.\) The number of iterations required in this case is \(k \sim \frac{\pi \sqrt{N}}{4},\) making the complexity of the algorithm \(\mathcal{O}(\sqrt{N}).\) This provides a quadratic speedup compared to the classical \(\mathcal{O}(N)\) brute-force approach.

Amplitude Amplification in PennyLane

Let’s take a look at a practical example using PennyLane: solving the zero-sum problem. In this task, we are given a list of \(n\) integers and our goal is to find all subsets of numbers whose sum is equal \(0.\) In this example, we use the following set of integers:

values = [1, -2, 3, 4, 5, -6]
n = len(values)

Can you find all the subsets that add to zero? 🤔

We will use Amplitude Amplification to solve the problem. First we define a binary variable \(x_i\) that takes the value \(1\) if we include the \(i\)-th element in the subset and takes the value \(0\) otherwise. We encode the \(i\)-th variable in the \(i\)-th qubit of a quantum state. For instance, \(|110001\rangle\) represents the subset \([1,-2,-6]\) consisting of the first, second, and sixth element in the set. Later on, we will see how to solve it directly with AmplitudeAmplification, but it is worthwhile to go step by step showing each part of the algorithm.

We can now define the initial state:

This is a uniform superposition of all possible subsets so solutions are guaranteed to have non-zero amplitudes in \(|\Psi\rangle.\) Let’s generate the state and visualize it.

import pennylane as qml
import matplotlib.pyplot as plt
plt.style.use('pennylane.drawer.plot')

@qml.prod
# This decorator converts the quantum function U into an operator.
# It is useful for later use in the AmplitudeAmplification template.
def U(wires):
    # create the generator U, such that U|0⟩ = |Ψ⟩
    for wire in wires:
        qml.Hadamard(wires=wire)

dev = qml.device("default.qubit")

@qml.qnode(dev)
def circuit():
    U(wires = range(n))
    return qml.state()

output = circuit()[:64].real

plt.bar(range(len(output)), output)
plt.ylim(-0.4, 0.6)
plt.ylabel("Amplitude")
plt.xlabel("|i⟩")
plt.axhline(0, color='black', linewidth=1)
plt.show()

Initially, the probability of getting any basis state is the same. The next step is to define the oracle that marks the elements that satisfy our property — that the sum of the subset is zero. To do this we first create an operator that stores the sum of the selected subset in some auxiliary qubits. This operator, which we call \(\text{Sum}\) , is defined as:

where \(v_i\) is the \(i\)-th integer in the input set. For more details of how we build this operation take a look at Basic arithmetic with the Quantum Fourier Transform.

import numpy as np

def Sum(wires_subset, wires_sum):
    qml.QFT(wires = wires_sum)
    for i, value in enumerate(values):
        for j in range(len(wires_sum)):
            qml.CRZ(value * np.pi / (2 ** j), wires=[wires_subset[i], wires_sum[j]])
    qml.adjoint(qml.QFT)(wires=wires_sum)

To create the oracle that performs the reflection around \(|\phi^{\perp}\rangle,\) we apply the \(\text{Sum}\) operator to the state and then flip the sign of those states whose sum is \(0.\) This allows us to mark the searched elements. Then we apply the inverse of the sum to clean the auxiliary qubits.

@qml.prod
def oracle(wires_subset, wires_sum):
    # Reflection on |ϕ⟂⟩
    Sum(wires_subset, wires_sum)
    qml.FlipSign(0, wires=wires_sum)
    qml.adjoint(Sum)(wires_subset, wires_sum)


@qml.qnode(dev)
def circuit():
    U(wires=range(n))                 # Generate initial state
    oracle(range(n), range(n, n+5))   # Apply the reflection on |ϕ⟂⟩
    return qml.state()


output = circuit()[0::2 ** 5].real
plt.bar(range(len(output)), output)
plt.ylim(-0.4, 0.6)
plt.ylabel("Amplitude")
plt.xlabel("|i⟩")
plt.axhline(0, color='black', linewidth=1)
plt.show()

Great, we have flipped the sign of the solution states. Note that in this example we can check every possible subset explicitly, but this becomes exponentially hard for larger input sets. Still, this operator would correctly flip the sign of any solution. The next step is to reflect on the \(|\Psi\rangle\) state. This reflection operator can be built directly in PennyLane with Reflection. The final circuit is then:

@qml.qnode(dev)
def circuit():
    U(wires=range(n))                  # Generate initial state
    oracle(range(n), range(n, n + 5))  # Apply the reflection on |ϕ⟂⟩
    qml.Reflection(U(wires=range(n)))  # Reflect on |Ψ⟩
    return qml.state()

Let’s now look at the state \(|\Psi\rangle\) and see how it is changed.

output = circuit()[0::2 ** 5].real
plt.bar(range(len(output)), output)
plt.ylim(-0.4, 0.6)
plt.ylabel("Amplitude")
plt.xlabel("|i⟩")
plt.axhline(0, color='black', linewidth=1)
plt.show()

We have now amplified the amplitude of all the states that represent a solution to our problem. The four peaks are obtained in \(0\), \(27\), \(35\) and \(61,\) whose binary representation corresponds with \(|000000\rangle\), \(|011011\rangle,\) \(|100011\rangle\) and \(|111101\rangle\) respectively. These states satisfy the property that the sum of the subset is \(0.\)

Let’s use the AmplitudeAmplification to code the problem in a more compact way. This template expects as input, the unitary \(U\) that prepares \(|\Psi\rangle,\) the reflection with respect to \(|\phi^{\perp}\rangle\) (i.e. the oracle), and the number of iterations. We increase the number of iterations in order to study the evolution of the initial state:

@qml.qnode(dev)
def circuit(iters):
    U(wires=range(n))
    qml.AmplitudeAmplification(U = U(wires = range(n)),
                               O = oracle(range(n), range(n, n + 5)),
                               iters = iters)

    return qml.probs(wires = range(n))

fig, axs = plt.subplots(2, 2, figsize=(14, 10))
for i in range(1,9,2):
    output = circuit(iters=i)
    ax = axs[i // 4, i //2 % 2]
    ax.bar(range(len(output)), output)
    ax.set_ylim(0, 0.6)
    ax.set_title(f"Iteration {i}")

plt.tight_layout()
plt.subplots_adjust(bottom=0.1)
plt.axhline(0, color='black', linewidth=1)
plt.show()

We can see that as the number of iterations increases, the probability of success increases as well. But we should be careful to not overdo it: after 5 iterations in our example, the amplitudes have in fact decreased. This phenomenon is known as “overcooking” and is a consequence of rotating the state too much. It can be addressed using fixed-point amplitude amplification, which we discuss below.