barren_plateaus.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.

Barren plateaus

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver fuente en GitHub

Descargar el bloc de notas

En este ejemplo analizaremos los resultados de McClean, 2019 que dice que no cualquier estructura de red neuronal cuántica tendrá un buen desempeño para el aprendizaje. En particular, verá que una familia grande cualquiera de circuitos cuánticos aleatorios no sirve como una red neuronal cuántica buena, ya que tiene gradientes que se desvanecen casi en todas partes. En este ejemplo, no entrenaremos ningún modelo para un problema de aprendizaje específico, sino que nos centraremos en un problema más simple, el de entender el comportamiento de los gradientes.

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

np.random.seed(1234)

1. Resumen

Los circuitos cuánticos aleatorios con muchos bloques que tienen el siguiente aspecto ($R_{P}(\theta)$ es una rotación Pauli):

Donde si $f(x)$ se define como el valor de expectativa w.r.t. $Z_{a}Z_{b}$ para cualquier bit cuántico $a$ y $b$, entonces, hay un problema, que la media de $f'(x)$ es muy cercana a 0 y no varía mucho más. Lo notaremos a continuación:

2. Generación de circuitos cuánticos

La construcción de la publicación es fácil de seguir. Se implementa una función simple que genera un circuito cuántico aleatorio, a veces denominado red neuronal cuántica (QNN), con la profundidad dada en un conjunto de bits cuánticos:

def generate_random_qnn(qubits, symbol, depth):
    """Generate random QNN's with the same structure from McClean et al."""
    circuit = cirq.Circuit()
    for qubit in qubits:
        circuit += cirq.ry(np.pi / 4.0)(qubit)

    for d in range(depth):
        # Add a series of single qubit rotations.
        for i, qubit in enumerate(qubits):
            random_n = np.random.uniform()
            random_rot = np.random.uniform(
            ) * 2.0 * np.pi if i != 0 or d != 0 else symbol
            if random_n > 2. / 3.:
                # Add a Z.
                circuit += cirq.rz(random_rot)(qubit)
            elif random_n > 1. / 3.:
                # Add a Y.
                circuit += cirq.ry(random_rot)(qubit)
            else:
                # Add a X.
                circuit += cirq.rx(random_rot)(qubit)

        # Add CZ ladder.
        for src, dest in zip(qubits, qubits[1:]):
            circuit += cirq.CZ(src, dest)

    return circuit


generate_random_qnn(cirq.GridQubit.rect(1, 3), sympy.Symbol('theta'), 2)

Los autores investigan el gradiente de un solo parámetro $\theta_{1,1}$. Sigamos y coloquemos un sympy.Symbol en el circuito donde estaría $\theta_{1,1}$. Como los autores no analizan las estadísticas de ningún otro símbolo del circuito, reemplacémoslo con valores aleatorios ahora, en vez de más tarde.

3. Ejecución de los circuitos

Generemos algunos de estos circuitos con un observable para hacer pruebas de la afirmación de que los gradientes no varían mucho. Primero, generemos un lote de circuitos aleatorios. Elijamos un observable ZZ aleatorio y calculemos el lote de gradientes y la variancia con TensorFlow Quantum.

3.1 Cálculo de la variancia por lotes

Escribamos una función ayudante que calcule la variancia del gradiente de un observable sobre un conjunto de circuitos:

def process_batch(circuits, symbol, op):
    """Compute the variance of a batch of expectations w.r.t. op on each circuit that 
    contains `symbol`. Note that this method sets up a new compute graph every time it is
    called so it isn't as performant as possible."""

    # Setup a simple layer to batch compute the expectation gradients.
    expectation = tfq.layers.Expectation()

    # Prep the inputs as tensors
    circuit_tensor = tfq.convert_to_tensor(circuits)
    values_tensor = tf.convert_to_tensor(
        np.random.uniform(0, 2 * np.pi, (n_circuits, 1)).astype(np.float32))

    # Use TensorFlow GradientTape to track gradients.
    with tf.GradientTape() as g:
        g.watch(values_tensor)
        forward = expectation(circuit_tensor,
                              operators=op,
                              symbol_names=[symbol],
                              symbol_values=values_tensor)

    # Return variance of gradients across all circuits.
    grads = g.gradient(forward, values_tensor)
    grad_var = tf.math.reduce_std(grads, axis=0)
    return grad_var.numpy()[0]

3.1 Preparación y ejecución

Elijamos la cantidad de circuitos aleatorios que generaremos junto con su profundidad y la cantidad de bit cuánticos que deberían participar. Después, grafiquemos los resultados.

n_qubits = [2 * i for i in range(2, 7)
           ]  # Ranges studied in paper are between 2 and 24.
depth = 50  # Ranges studied in paper are between 50 and 500.
n_circuits = 200
theta_var = []

for n in n_qubits:
    # Generate the random circuits and observable for the given n.
    qubits = cirq.GridQubit.rect(1, n)
    symbol = sympy.Symbol('theta')
    circuits = [
        generate_random_qnn(qubits, symbol, depth) for _ in range(n_circuits)
    ]
    op = cirq.Z(qubits[0]) * cirq.Z(qubits[1])
    theta_var.append(process_batch(circuits, symbol, op))

plt.semilogy(n_qubits, theta_var)
plt.title('Gradient Variance in QNNs')
plt.xlabel('n_qubits')
plt.xticks(n_qubits)
plt.ylabel('$\\partial \\theta$ variance')
plt.show()

En este gráfico se muestra que para los problemas de aprendizaje automático cuántico, no se puede, simplemente, adivinar un ansatz de QNN aleatorio y esperar obtener el mejor resultado. En el circuito del modelo, debe haber alguna estructura presente para que los gradientes varíen al punto en que el aprendizaje efectivamente se produzca.

4. Heurística

Una heurística interesante de Grant, 2019 permite empezar muy cerca de la opción aleatoria, pero tampoco tanto. Con los mismos circuitos que McClean et al., los autores proponen una técnica de inicialización diferente para los parámetros de control clásico para evitar los barren plateaus. La técnica de inicialización inicia algunas capas con parámetros de control totalmente aleatorios. Pero en las capas que la siguen inmediatamente, elige parámetros tales que la transformación inicial hecha por las primeras pocas capas se deshace. A esto los autores lo llaman bloque de identidad.

La ventaja de esta heurística es que al cambiar un solo parámetro, todos los otros bloques que se encuentran fuera del bloque actual seguirán siendo de identidad; y la señal de gradiente consigue volverse mucho más fuerte que antes. De este modo, el usuario puede elegir qué variables y bloques desea modificar para obtener una señal de gradiente más fuerte. Esta heurística no impide que el usuario caiga en un barren plateau durante la fase de entrenamiento (y restringe una actualización totalmente simultanea), simplemente, garantiza que se podrá comenzar fuera del plateau.

4.1 Construcción de una nueva red neuronal cuántica

Ahora, construimos una función para generar redes neuronales cuánticas de bloques de identidad. Esta implementación es un poco diferente a la de la publicación. Por ahora, observamos el comportamiento del gradiente de un único parámetro para ser coherentes con McClean et al, y para que se puedan hacer algunas simplificaciones.

Para generar un bloque de identidad y entrenar el modelo, por lo general necesitamos $U1(\theta_{1a}) U1(\theta_{1b})^{\dagger}$ y no $U1(\theta_1) U1(\theta_1)^{\dagger}$. Inicialmente $\theta_{1a}$ y $\theta_{1b}$ son los mismos ángulos, pero se aprenden de forma diferente. De lo contrario, deberíamos obtener siempre la identidad, incluso después del entrenamiento. La elección de la cantidad de bloques de identidad es empírica. Mientras más profundo sea el bloque, más pequeña será la variancia en medio de ese bloque. Pero al principio y al final del bloque, la variancia de los gradientes de parámetro debería ser grande.

def generate_identity_qnn(qubits, symbol, block_depth, total_depth):
    """Generate random QNN's with the same structure from Grant et al."""
    circuit = cirq.Circuit()

    # Generate initial block with symbol.
    prep_and_U = generate_random_qnn(qubits, symbol, block_depth)
    circuit += prep_and_U

    # Generate dagger of initial block without symbol.
    U_dagger = (prep_and_U[1:])**-1
    circuit += cirq.resolve_parameters(
        U_dagger, param_resolver={symbol: np.random.uniform() * 2 * np.pi})

    for d in range(total_depth - 1):
        # Get a random QNN.
        prep_and_U_circuit = generate_random_qnn(
            qubits,
            np.random.uniform() * 2 * np.pi, block_depth)

        # Remove the state-prep component
        U_circuit = prep_and_U_circuit[1:]

        # Add U
        circuit += U_circuit

        # Add U^dagger
        circuit += U_circuit**-1

    return circuit


generate_identity_qnn(cirq.GridQubit.rect(1, 3), sympy.Symbol('theta'), 2, 2)

4.2 Comparación

Aquí podemos ver que la heurística ayuda a mantener la variancia del gradiente para que no se desvanezca, ya que rápidamente, hace lo siguiente:

block_depth = 10
total_depth = 5

heuristic_theta_var = []

for n in n_qubits:
    # Generate the identity block circuits and observable for the given n.
    qubits = cirq.GridQubit.rect(1, n)
    symbol = sympy.Symbol('theta')
    circuits = [
        generate_identity_qnn(qubits, symbol, block_depth, total_depth)
        for _ in range(n_circuits)
    ]
    op = cirq.Z(qubits[0]) * cirq.Z(qubits[1])
    heuristic_theta_var.append(process_batch(circuits, symbol, op))

plt.semilogy(n_qubits, theta_var)
plt.semilogy(n_qubits, heuristic_theta_var)
plt.title('Heuristic vs. Random')
plt.xlabel('n_qubits')
plt.xticks(n_qubits)
plt.ylabel('$\\partial \\theta$ variance')
plt.show()

Es una excelente mejora para obtener señales de gradientes más fuertes a partir de redes neuronales cuánticas aleatorias (cercanas).