gradients.ipynb

Kernel: Python 3

Copyright 2020 The TensorFlow Authors.

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Cálculo de gradientes

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver fuente en GitHub

Descargar el bloc de notas

En este tutorial se exploran los algoritmos con los que se calculan los gradientes para los valores esperados de los circuitos cuánticos.

El cálculo del gradiente de un valor esperado de un observable en un circuito cuántico es un proceso complejo. Los valores esperados de los observables no se pueden dar el lujo de tener fórmulas de gradientes analíticos que siempre sean fáciles de escribir; a diferencia de lo que sucede con las transformaciones de aprendizaje automático tradicional, como la multiplicación matricial o la suma de vectores que tienen fórmulas de gradientes analíticos que son fáciles de escribir. Como resultado, hay diferentes métodos para calcular los gradientes cuánticos que resultan prácticos para los diferentes escenarios posibles. En este tutorial se comparan y contrastan dos esquemas de diferenciación distintos.

Preparación

!pip install tensorflow==2.7.0

Instalar TensorFlow Quantum:

!pip install tensorflow-quantum==0.7.2

# Update package resources to account for version changes.
import importlib, pkg_resources
importlib.reload(pkg_resources)

Ahora, hay que importar TensorFlow y las dependencias del módulo:

import tensorflow as tf
import tensorflow_quantum as tfq

import cirq
import sympy
import numpy as np

# visualization tools
%matplotlib inline
import matplotlib.pyplot as plt
from cirq.contrib.svg import SVGCircuit

1. Avance preliminar

Formémonos una idea algo más concreta de lo que es el cálculo de gradientes para circuitos cuánticos. Supongamos que tenemos un circuito parametrizado como el siguiente:

qubit = cirq.GridQubit(0, 0)
my_circuit = cirq.Circuit(cirq.Y(qubit)**sympy.Symbol('alpha'))
SVGCircuit(my_circuit)

Junto con un observable:

pauli_x = cirq.X(qubit)
pauli_x

Al observar este operador sabemos que $⟨Y(\alpha)| X | Y(\alpha)⟩ = \sin(\pi \alpha)$

def my_expectation(op, alpha):
    """Compute ⟨Y(alpha)| `op` | Y(alpha)⟩"""
    params = {'alpha': alpha}
    sim = cirq.Simulator()
    final_state_vector = sim.simulate(my_circuit, params).final_state_vector
    return op.expectation_from_state_vector(final_state_vector, {qubit: 0}).real


my_alpha = 0.3
print("Expectation=", my_expectation(pauli_x, my_alpha))
print("Sin Formula=", np.sin(np.pi * my_alpha))

y si definimos $f_{1}(\alpha) = ⟨Y(\alpha)| X | Y(\alpha)⟩$ entonces $f_{1}^{'}(\alpha) = \pi \cos(\pi \alpha)$. Veamos esto:

def my_grad(obs, alpha, eps=0.01):
    grad = 0
    f_x = my_expectation(obs, alpha)
    f_x_prime = my_expectation(obs, alpha + eps)
    return ((f_x_prime - f_x) / eps).real


print('Finite difference:', my_grad(pauli_x, my_alpha))
print('Cosine formula:   ', np.pi * np.cos(np.pi * my_alpha))

2. La necesidad de un diferenciador

Con circuitos más grandes, no siempre uno tiene la suerte de contar con una fórmula que calcule con precisión los gradientes de un circuito cuántico dado. En caso de que una fórmula simple no sea suficiente para calcular el gradiente, la clase tfq.differentiators.Differentiator permite definir algoritmos para calcular los gradientes de los circuitos. Por ejemplo, podemos recrear el ejemplo anterior en TensorFlow Quantum (TFQ) con:

expectation_calculation = tfq.layers.Expectation(
    differentiator=tfq.differentiators.ForwardDifference(grid_spacing=0.01))

expectation_calculation(my_circuit,
                        operators=pauli_x,
                        symbol_names=['alpha'],
                        symbol_values=[[my_alpha]])

Sin embargo, si cambiamos y estimamos la esperanza basándonos en las muestras (lo que pasaría en un dispositivo verdadero) los valores podrían variar un poco. Lo que significa que ahora tenemos una estimación imperfecta:

sampled_expectation_calculation = tfq.layers.SampledExpectation(
    differentiator=tfq.differentiators.ForwardDifference(grid_spacing=0.01))

sampled_expectation_calculation(my_circuit,
                                operators=pauli_x,
                                repetitions=500,
                                symbol_names=['alpha'],
                                symbol_values=[[my_alpha]])

Esto se puede agravarse rápidamente y causar un problema grave con la exactitud en lo que respecta a los gradientes:

# Make input_points = [batch_size, 1] array.
input_points = np.linspace(0, 5, 200)[:, np.newaxis].astype(np.float32)
exact_outputs = expectation_calculation(my_circuit,
                                        operators=pauli_x,
                                        symbol_names=['alpha'],
                                        symbol_values=input_points)
imperfect_outputs = sampled_expectation_calculation(my_circuit,
                                                    operators=pauli_x,
                                                    repetitions=500,
                                                    symbol_names=['alpha'],
                                                    symbol_values=input_points)
plt.title('Forward Pass Values')
plt.xlabel('$x$')
plt.ylabel('$f(x)$')
plt.plot(input_points, exact_outputs, label='Analytic')
plt.plot(input_points, imperfect_outputs, label='Sampled')
plt.legend()

# Gradients are a much different story.
values_tensor = tf.convert_to_tensor(input_points)

with tf.GradientTape() as g:
    g.watch(values_tensor)
    exact_outputs = expectation_calculation(my_circuit,
                                            operators=pauli_x,
                                            symbol_names=['alpha'],
                                            symbol_values=values_tensor)
analytic_finite_diff_gradients = g.gradient(exact_outputs, values_tensor)

with tf.GradientTape() as g:
    g.watch(values_tensor)
    imperfect_outputs = sampled_expectation_calculation(
        my_circuit,
        operators=pauli_x,
        repetitions=500,
        symbol_names=['alpha'],
        symbol_values=values_tensor)
sampled_finite_diff_gradients = g.gradient(imperfect_outputs, values_tensor)

plt.title('Gradient Values')
plt.xlabel('$x$')
plt.ylabel('$f^{\'}(x)$')
plt.plot(input_points, analytic_finite_diff_gradients, label='Analytic')
plt.plot(input_points, sampled_finite_diff_gradients, label='Sampled')
plt.legend()

Aquí se puede observar que, a pesar de que la fórmula de diferencia finita funciona rápido para calcular los gradientes en un caso analítico, cuando trabajamos con métodos basados en muestras se vuelve demasiado ruidosa. Por lo tanto, conviene utilizar técnicas más minuciosas para garantizar que se calcule un buen gradiente. A continuación, observaremos una técnica mucho más lenta que probablemente tampoco sea la ideal para los cálculos de gradientes de esperanza analítica, pero que tenga un desempeño superior en un caso basado en una muestra del mundo real:

# A smarter differentiation scheme.
gradient_safe_sampled_expectation = tfq.layers.SampledExpectation(
    differentiator=tfq.differentiators.ParameterShift())

with tf.GradientTape() as g:
    g.watch(values_tensor)
    imperfect_outputs = gradient_safe_sampled_expectation(
        my_circuit,
        operators=pauli_x,
        repetitions=500,
        symbol_names=['alpha'],
        symbol_values=values_tensor)

sampled_param_shift_gradients = g.gradient(imperfect_outputs, values_tensor)

plt.title('Gradient Values')
plt.xlabel('$x$')
plt.ylabel('$f^{\'}(x)$')
plt.plot(input_points, analytic_finite_diff_gradients, label='Analytic')
plt.plot(input_points, sampled_param_shift_gradients, label='Sampled')
plt.legend()

A partir de lo anterior, podemos ver que existen ciertos diferenciadores que funcionan mejor con algunos escenarios de investigación en particular. En general, los métodos basados en muestras más lentos que son robustos ante el ruido del dispositivo, entre otras cosas, son excelentes diferenciadores, cuando se usan en pruebas o se implementan algoritmos en un entorno más cercano al "mundo real". Los métodos más rápidos como los de diferencia finita son excelentes para los cálculos analíticos y cuando se desea obtener un mayor rendimiento, pero no están vinculados aún a la viabilidad del algoritmo en el dispositivo.

3. Muchos observables

Presentemos un segundo observable y veamos la compatibilidad que ofrece TensorFlow Quantum a múltiples observables para un solo circuito.

pauli_z = cirq.Z(qubit)
pauli_z

Si este observable se usa con el mismo circuito de antes, entonces, tenemos $f_{2}(\alpha) = ⟨Y(\alpha)| Z | Y(\alpha)⟩ = \cos(\pi \alpha)$ y $f_{2}^{'}(\alpha) = -\pi \sin(\pi \alpha)$. Hagamos una comprobación rápida:

test_value = 0.

print('Finite difference:', my_grad(pauli_z, test_value))
print('Sin formula:      ', -np.pi * np.sin(np.pi * test_value))

Coincide (o está lo suficientemente cerca).

Ahora, si definimos $g(\alpha) = f_{1}(\alpha) + f_{2}(\alpha)$, entonces $g'(\alpha) = f_{1}^{'}(\alpha) + f^{'}_{2}(\alpha)$. La determinación de más de un observable en TensorFlow Quantum para usar junto con un circuito equivale a agregar más términos a $g$.

Significa que el gradiente de un símbolo en particular dentro de un circuito es igual a la suma de los gradientes respecto de cada observable de ese símbolo aplicado a ese mismo circuito. Esto es compatible con la toma y propagación hacia atrás de gradientes de TensorFlow (dada la suma de los gradientes de todos los observables como el gradiente de un símbolo en particular).

sum_of_outputs = tfq.layers.Expectation(
    differentiator=tfq.differentiators.ForwardDifference(grid_spacing=0.01))

sum_of_outputs(my_circuit,
               operators=[pauli_x, pauli_z],
               symbol_names=['alpha'],
               symbol_values=[[test_value]])

Aquí observamos que la primera entrada es la esperanza de w.r.t Pauli X y la segunda es la esperanza de w.r.t Pauli Z. Ahora, cuando tomamos el gradiente:

test_value_tensor = tf.convert_to_tensor([[test_value]])

with tf.GradientTape() as g:
    g.watch(test_value_tensor)
    outputs = sum_of_outputs(my_circuit,
                             operators=[pauli_x, pauli_z],
                             symbol_names=['alpha'],
                             symbol_values=test_value_tensor)

sum_of_gradients = g.gradient(outputs, test_value_tensor)

print(my_grad(pauli_x, test_value) + my_grad(pauli_z, test_value))
print(sum_of_gradients.numpy())

Hemos verificado que la suma de los gradientes para cada observable es, en definitiva, el gradiente de $\alpha$. Este comportamiento es compatible con todos los diferenciadores de TensorFlow Quantum y es crucial para la compatibilidad con el resto de TensorFlow.

4. Uso avanzado

Todos los diferenciadores que existen dentro de TensorFlow Quantum tienen la subclase tfq.differentiators.Differentiator. Para implementar un diferenciador, un usuario debe implementar una de dos interfaces. Lo estándar es implementar get_gradient_circuits, que le indica a la clase de base qué circuitos medir para obtener una estimación del gradiente. Como alternativa, se pueden sobrecargar differentiate_analytic y differentiate_sampled; la clase tfq.differentiators.Adjoint toma su ruta.

En el siguiente caso se usa TensorFlow Quantum para implementar el gradiente de un circuito. Usaremos un ejemplo pequeño de un cambio de parámetro.

Retomemos el circuito definido arriba en este texto, $|\alpha⟩ = Y^{\alpha}|0⟩$. Al igual que antes, podemos definir una función como el valor esperado de este circuito con respecto al observable $X$, $f(\alpha) = ⟨\alpha|X|\alpha⟩$. Al utilizar las reglas para cambios de parámetros, en este circuito, podemos hallar que la derivada es $$\frac{\partial}{\partial \alpha} f(\alpha) = \frac{\pi}{2} f\left(\alpha + \frac{1}{2}\right) - \frac{ \pi}{2} f\left(\alpha - \frac{1}{2}\right)$$. La función get_gradient_circuits devuelve los componentes de esta derivada.

class MyDifferentiator(tfq.differentiators.Differentiator):
    """A Toy differentiator for <Y^alpha | X |Y^alpha>."""

    def __init__(self):
        pass

    def get_gradient_circuits(self, programs, symbol_names, symbol_values):
        """Return circuits to compute gradients for given forward pass circuits.
        
        Every gradient on a quantum computer can be computed via measurements
        of transformed quantum circuits.  Here, you implement a custom gradient
        for a specific circuit.  For a real differentiator, you will need to
        implement this function in a more general way.  See the differentiator
        implementations in the TFQ library for examples.
        """

        # The two terms in the derivative are the same circuit...
        batch_programs = tf.stack([programs, programs], axis=1)

        # ... with shifted parameter values.
        shift = tf.constant(1/2)
        forward = symbol_values + shift
        backward = symbol_values - shift
        batch_symbol_values = tf.stack([forward, backward], axis=1)
  
        # Weights are the coefficients of the terms in the derivative.
        num_program_copies = tf.shape(batch_programs)[0]
        batch_weights = tf.tile(tf.constant([[[np.pi/2, -np.pi/2]]]),
                                [num_program_copies, 1, 1])

        # The index map simply says which weights go with which circuits.
        batch_mapper = tf.tile(
            tf.constant([[[0, 1]]]), [num_program_copies, 1, 1])

        return (batch_programs, symbol_names, batch_symbol_values,
                batch_weights, batch_mapper)

La clase de base Differentiator usa los componentes devueltos por get_gradient_circuits para calcular la derivada, como en la fórmula para el cambio de parámetro que vimos arriba. Este nuevo diferenciador, ahora, puede usarse con objetos tfq.layer existentes:

custom_dif = MyDifferentiator()
custom_grad_expectation = tfq.layers.Expectation(differentiator=custom_dif)

# Now let's get the gradients with finite diff.
with tf.GradientTape() as g:
    g.watch(values_tensor)
    exact_outputs = expectation_calculation(my_circuit,
                                            operators=[pauli_x],
                                            symbol_names=['alpha'],
                                            symbol_values=values_tensor)

analytic_finite_diff_gradients = g.gradient(exact_outputs, values_tensor)

# Now let's get the gradients with custom diff.
with tf.GradientTape() as g:
    g.watch(values_tensor)
    my_outputs = custom_grad_expectation(my_circuit,
                                         operators=[pauli_x],
                                         symbol_names=['alpha'],
                                         symbol_values=values_tensor)

my_gradients = g.gradient(my_outputs, values_tensor)

plt.subplot(1, 2, 1)
plt.title('Exact Gradient')
plt.plot(input_points, analytic_finite_diff_gradients.numpy())
plt.xlabel('x')
plt.ylabel('f(x)')
plt.subplot(1, 2, 2)
plt.title('My Gradient')
plt.plot(input_points, my_gradients.numpy())
plt.xlabel('x')

Este nuevo diferenciador, ahora se puede usar para generar operaciones diferenciables.

Punto clave: Si un diferenciador que ya se ha adjuntado previamente a una operación, antes de adjuntarlo a la operación nueva, hay que actualizarlo. Porque un diferenciador solamente debe adjuntarse a una operación por vez.

# Create a noisy sample based expectation op.
expectation_sampled = tfq.get_sampled_expectation_op(
    cirq.DensityMatrixSimulator(noise=cirq.depolarize(0.01)))

# Make it differentiable with your differentiator:
# Remember to refresh the differentiator before attaching the new op
custom_dif.refresh()
differentiable_op = custom_dif.generate_differentiable_op(
    sampled_op=expectation_sampled)

# Prep op inputs.
circuit_tensor = tfq.convert_to_tensor([my_circuit])
op_tensor = tfq.convert_to_tensor([[pauli_x]])
single_value = tf.convert_to_tensor([[my_alpha]])
num_samples_tensor = tf.convert_to_tensor([[5000]])

with tf.GradientTape() as g:
    g.watch(single_value)
    forward_output = differentiable_op(circuit_tensor, ['alpha'], single_value,
                                       op_tensor, num_samples_tensor)

my_gradients = g.gradient(forward_output, single_value)

print('---TFQ---')
print('Foward:  ', forward_output.numpy())
print('Gradient:', my_gradients.numpy())
print('---Original---')
print('Forward: ', my_expectation(pauli_x, my_alpha))
print('Gradient:', my_grad(pauli_x, my_alpha))

Excelente: Ahora, podemos usar todos los diferenciadores que TensorFlow Quantum tiene para ofrecer y, además, podemos definir el propio.