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 em TensorFlow.org

Executar no Google Colab

Ver fonte no GitHub

Baixar notebook

Neste exemplo, você explorará o resultado de McClean, 2019, que afirma que nem toda estrutura de rede neural quântica terá um bom desempenho quando se trata do aprendizado. Em particular, você verá que uma certa grande família de circuitos quânticos aleatórios não serve como boas redes neurais quânticas, porque têm gradientes que desaparecem em quase todo lugar. Neste exemplo, você não treinará nenhum modelo para um problema de aprendizado específico, e sim focará no problema mais simples de entender o comportamento dos gradientes.

Configuração

!pip install tensorflow==2.7.0

Instale o 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)

Agora importe o TensorFlow e as dependências de 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. Resumo

Circuitos quânticos aleatórios com vários blocos que se parecem assim ($R_{P}(\theta)$ é uma rotação de Pauli aleatória):

Em que $f(x)$ é definido como o valor esperado w.r.t. $Z_{a}Z_{b}$ para quaisquer qubits $a$ e $b$, e também há o problema de $f'(x)$ ser uma média bastante próxima a 0 e não variar muito. Você verá isso abaixo:

2. Geração de circuitos aleatórios

É fácil seguir a construção do artigo. O seguinte código implementa uma função simples que gera um circuito quântico aleatório — às vezes chamado de rede neural quântica (QNN)— com a determinada profundidade em um conjunto de qubits:

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)

Os autores investigam o gradiente de um único parâmetro $\theta_{1,1}$. Vamos seguir isso ao colocar um sympy.Symbol no circuito, onde $\theta_{1,1}$ estaria. Como os autores não analisam as estatísticas de qualquer outro símbolo no circuito, vamos substituir por valores aleatórios agora, e não depois.

3. Execução dos circuitos

Gere alguns desses circuitos com um observável para testar a afirmação de que os gradientes não variam muito, Primeiro, gere um lote de circuitos aleatórios. Escolha um observável ZZ aleatório e calcule os gradientes e a variância em lote usando o TensorFlow Quantum.

3.1 Computação da variância em lote

Vamos escrever uma função helper que computa a variância do gradiente de um determinado observável em um lote 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 Configure e execute

Escolha o número de circuitos aleatórios para gerar com a profundidade deles e o número de qubits em que devem agir. Em seguida, plote os 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()

Essa plotagem mostra que, para problemas de aprendizado de máquina quântico, você não pode simplesmente chutar um ansatz de QNN aleatório e torcer pelo melhor. O circuito do modelo precisa ter alguma estrutura para que os gradientes variem até o ponto em que o aprendizado é possível.

4. Heurística

Uma heurística interessante por Grant, 2019 permite uma inicialização bem próxima do aleatório, mas não exatamente. Usando os mesmos circuitos que McClean et al., os autores propõem uma técnica de inicialização diferente para os parâmetros de controle clássicos evitarem barren plateaus. Nessa técnica, são iniciadas algumas camadas com parâmetros de controle totalmente aleatórios — mas, nas camadas logo após, são escolhidos parâmetros para que essa transformação inicial das primeiras camadas seja desfeita. Os autores chamam isso de bloco de identidade.

A vantagem dessa heurística é que, ao mudar só um parâmetro, todos os outros blocos fora do atual permanecerão a identidade — e o sinal do gradiente chegará muito mais forte do que antes. Isso permite que o usuário escolha quais variáveis e blocos modificar para obter um forte sinal de gradiente. Essa heurística não evita que os usuários caiam em um barren plateau durante a fase de treinamento (e restringe uma atualização totalmente simultânea), só garante o começo fora de um plateau.

4.1 Nova construção de QNN

Agora construa uma função para gerar QNNs de bloco de identidade. Essa implementação é um pouco diferente daquela no artigo. Por enquanto, observe o comportamento do gradiente de um único parâmetro para que seja consistente com McClean et al. Assim, é possível fazer algumas simplificações.

Para gerar um bloco de identidade e treinar o modelo, geralmente, você precisa de $U1(\theta_{1a}) U1(\theta_{1b})^{\dagger}$, e não $U1(\theta_1) U1(\theta_1)^{\dagger}$. Inicialmente, $\theta_{1a}$ e $\theta_{1b}$ são os mesmos ângulos, mas eles são aprendidos de forma independente. Caso contrário, você sempre obterá a identidade, mesmo depois do treinamento. A escolha do número de blocos de identidade é empírica. Quanto mais profundo o bloco, menor a variância no meio dele. No entanto, no início e no final do bloco, a variância dos gradientes dos parâmetros deve 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 Comparação

Aqui você pode ver que a heurística realmente ajuda a impedir que a variância do gradiente desapareça rapidamente:

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()

Essa é uma ótima melhoria para receber sinais de gradiente mais fortes de QNNs (quase) aleatórias.