TensorFlow QuantumでQAOAをやってみる

はじめに

先日発表されたTensorFlow Quantumを触ってみました。
TensorFlowにCirqを統合し、さらに量子回路のパラメーターをニューラルネットで調整する機能をKeras風味で実装した量子機械学習用ライブラリです。

公式Tutorialの構成がハードボイルドで、最初に(おそらく実機Qubitを想定した)ゲート回転角誤差の補正計算から入り、次にVQEやQAOAなどのNear term algorithmをすっ飛ばしてMNISTを実装し、結果として現時点では古典NNの方がずっと優れていると結論付けています。

Near term algorithm との親和性が気になるところなので、今回はTensorFlow QuantumでQAOAを試してみます。

また今回の記事内ではQAOAにニューラルネットワークを使うことでパフォーマンスが向上するかといった議論はしませんし、できないと思います。
ニューラルネットワーク構造やハイパーパラメータは特に追い込んでいないため、興味のある方はより良い結果が得られるか色々試すのも良いと思います。

題材

題材はこちらの最適化問題をとりあえずそのまま使いました。
公式Tutorial "hello_many_worlds"のコードをベースとしているので、合わせて見ることをおすすめします。

実装

ライブラリインポート

import tensorflow as tf
import tensorflow_quantum as tfq

import cirq
import sympy
import numpy as np
from cirq.contrib.svg import SVGCircuit

パラメータ設定

# 変更可能なパラメータ
p = 3 # QAOAのステップ数
K = 20 # 制約条件式のハイパーパラメータ
command_param = 0.1 # ニューラルネットの入力、多分何でも良い

 # Qubit数, 問題によって決まる
N = 6

# QAOA 回転角パラメータをsympy symbolとして用意
control_params = []
for i in range(p):
    gamma = 'gamma' + str(i)
    beta = 'beta' + str(i)
    control_params.append(sympy.symbols(gamma))
    control_params.append(sympy.symbols(beta))

QAOA用ユニタリゲートの用意

元記事の制約条件QUBO式をPauli-Z演算子で構成されたコストハミルトニアンに変換し、時間発展ゲートを実装しました。
ミキサハミルトニアン時間発展と初期状態準備用のゲートも用意しておきます。

# コストハミルトニアン時間発展
def CostUnitary(circuit, control_params, i, K):
    circuit.append(cirq.ZZ(qubits[0], qubits[1])**((K / 2 + 1) * control_params[i]))
    circuit.append(cirq.ZZ(qubits[0], qubits[2])**(K / 2 * control_params[i]))
    circuit.append(cirq.Z(qubits[0])**((-K / 2 - 6) * control_params[i]))
    circuit.append(cirq.ZZ(qubits[1], qubits[2])**((K + 1) / 2 * control_params[i]))
    circuit.append(cirq.Z(qubits[1])**((- K / 2 - 7) * control_params[i]))
    circuit.append(cirq.Z(qubits[2])**((- K / 2 - 5) * control_params[i]))
    circuit.append(cirq.ZZ(qubits[3], qubits[4])**((K / 2 + 1) * control_params[i]))
    circuit.append(cirq.ZZ(qubits[3], qubits[5])**(K / 2 * control_params[i]))
    circuit.append(cirq.Z(qubits[3])**((- K / 2 - 6) * control_params[i]))
    circuit.append(cirq.ZZ(qubits[4], qubits[5])**((K + 1) / 2 * control_params[i]))
    circuit.append(cirq.Z(qubits[4])**((- K / 2 - 7) * control_params[i]))
    circuit.append(cirq.Z(qubits[5])**((- K / 2 - 5) * control_params[i]))
    circuit.append(cirq.ZZ(qubits[0], qubits[3])**(2 * control_params[i]))
    circuit.append(cirq.ZZ(qubits[0], qubits[4])**(control_params[i]))
    circuit.append(cirq.ZZ(qubits[1], qubits[3])**(control_params[i]))
    circuit.append(cirq.ZZ(qubits[1], qubits[4])**(2 * control_params[i]))
    circuit.append(cirq.ZZ(qubits[1], qubits[5])**(1 / 2 * control_params[i]))
    circuit.append(cirq.ZZ(qubits[2], qubits[4])**(1 / 2 * control_params[i]))
    circuit.append(cirq.ZZ(qubits[2], qubits[5])**(2 * control_params[i]))
    return circuit

# ミキサハミルトニアン時間発展
def Mixer(circuit, i, N):
    for j in range(N):
        circuit.append(cirq.XPowGate(exponent = control_params[i]).on(qubits[j]))
    return circuit

# 初期状態(|+>)準備
def init_circuit(N):
    initCircuit = cirq.Circuit()
    qubits = cirq.GridQubit.rect(1, N)
    for j in range(N):
        initCircuit.append(cirq.H(qubits[j]))
    return initCircuit

量子回路生成

上で用意したゲートをQAOAステップ数$p$回作用させます。
init_circuit()は後で使います。

qubits = [cirq.GridQubit(i, 0) for i in range(N)]
model_circuit = cirq.Circuit()
for i in range(p):
    model_circuit = CostUnitary(model_circuit, control_params, 2 * i, K)
    model_circuit = Mixer(model_circuit, 2 * i + 1, N)

Operator定義

コストハミルトニアンの期待値を目的関数として最適化を行うため、コストハミルトニアンをOperatorとして実装します。

def ZZoperator(i, j):
    return cirq.Z(qubits[i]) * cirq.Z(qubits[j])

operators = [[-1 * ((K / 2 + 1) * ZZoperator(0, 1) + (K / 2) * ZZoperator(0, 2) - \
            (K / 2 + 6) * cirq.Z(qubits[0]) + (K / 2 + 1 / 2) * ZZoperator(1, 2) - \
            (K / 2 + 7) * cirq.Z(qubits[1]) - (K / 2 + 5) * cirq.Z(qubits[2]) + \
            (K / 2 + 1) * ZZoperator(3, 4) + (K / 2) * ZZoperator(3, 5) - \
            (K / 2 + 6) * cirq.Z(qubits[3]) + (K / 2 + 1 / 2) * ZZoperator(4, 5) - \
            (K / 2 + 7) * cirq.Z(qubits[4]) - (K / 2 + 5) * cirq.Z(qubits[5]) + \
            2 * ZZoperator(0, 3) + ZZoperator(0, 4) + ZZoperator(1, 3) + 2 * ZZoperator(1, 4) +\
            1 / 2 * ZZoperator(1, 5) + 1 / 2 * ZZoperator(2, 4) + 2 * ZZoperator(2, 5) + 2 * K + 24)]]

model定義

公式Tutorial "hello_many_worlds"を参考に学習用modelを定義していきます。
controllerには変分パラメータの値を出力するニューラルネットが入ります。
通常のkerasよろしくlayerを追加していきます。

commands_inputはcontrollerで組んだニューラルネットの先頭ノードへの入力です。
公式Tutorialでは複数の値を用意し、値によって異なる出力を返すようニューラルネットをトレーニングするなどしています。
今回は1つの値しか使わないため、あまり意味はないです。

最後にexpected_outputsを0としています。
この値と実際のコストハミルトニアンの期待値との差を小さくするようにニューラルネットのパラメータが調整されます。

expected_outputsはコストハミルトニアン期待値の大域的な最小値(最大値)としたいところですが、予めわかっていない方が一般的じゃないかと思います。
ここをどのような値にするか。今回は0未満の値をとらないQUBO式が与えられているためコストハミルトニアンの期待値も0に向かって最適化しました。

# QAOAパラメータ最適化用のネットワークを定義
controller = tf.keras.Sequential([
    tf.keras.layers.Dense(20 * p, activation='elu'),
    tf.keras.layers.Dense(len(control_params))
])

init_circuits = tfq.convert_to_tensor([init_circuit(N)])

commands_input = tf.keras.layers.Input(shape=(1),
                                       dtype=tf.dtypes.float32,
                                       name='commands_input')
circuits_input = tf.keras.Input(shape=(),
                                # The circuit-tensor has dtype `tf.string` 
                                dtype=tf.dtypes.string,
                                name='circuits_input')
operators_input = tf.keras.Input(shape=(1,),
                                 dtype=tf.dtypes.string,
                                 name='operators_input')

dense_2 = controller(commands_input) # tf.keras.Sequential

full_circuit = tfq.layers.AddCircuit()(circuits_input, append=model_circuit)

expectation_output = tfq.layers.Expectation()(full_circuit,
                                              symbol_names=control_params,
                                              symbol_values=dense_2,
                                              operators=operators_input)

model = tf.keras.Model(
    inputs=[circuits_input, commands_input, operators_input],
    outputs=[expectation_output])

operator_data = tfq.convert_to_tensor(operators)
commands = np.array([[command_param] for i in range(len(operators))], dtype=np.float32)
expected_outputs = np.array([[0] for i in range(len(operators))], dtype=np.float32)

modelコンパイル

optimizer = tf.keras.optimizers.Adam(learning_rate=0.05)
loss = tf.keras.losses.MeanSquaredError()

model.compile(optimizer=optimizer, loss=loss)

学習

history = model.fit(
    x=[init_circuits, commands, operator_data],
    y=expected_outputs,
    epochs=200,
    verbose=1)

コストハミルトニアン期待値、学習済パラメータの表示

print('Expectation value', model([init_circuits, commands, operator_data]))
after_params = controller.predict(np.array([command_param]))[0]
print('after params: ', after_params)

結果のサンプリング

ニューラルネットでパラメータ最適化された量子回路の出力を測定し、状態をサンプリングします。
得られた状態についてコストハミルトニアンの期待値を計算し、期待値最大の状態が今回の問題の最適解となります(元の制約条件式は最小化問題ですがコストハミルトニアンの符号を変えて最大化問題としています)。

param_dict = {}
for i in range(len(control_params)):
    param_dict.update([(control_params[i], after_params[i])])
resolver = cirq.ParamResolver(param_dict)

simulator = cirq.Simulator()
model_circuit.append(cirq.measure(*qubits, key = 'm'))
results = simulator.run(model_circuit, resolver, repetitions=100)

def calc_cost(state, operator):
    qubits = [cirq.GridQubit(i, 0) for i in range(len(state))]
    test_circuit = cirq.Circuit()
    for i in range(len(state)):
        if state[i] == '1':
            test_circuit.append(cirq.X(qubits[i]))
        else:
            test_circuit.append(cirq.I(qubits[i]))

    output_state_vector = cirq.Simulator().simulate(test_circuit).final_state
    qubit_map={qubits[0]: 0, qubits[1]: 1, qubits[2]: 2, qubits[3]: 3, qubits[4]: 4, qubits[5]: 5}
    return operator.expectation_from_wavefunction(output_state_vector, qubit_map).real

hist = results.histogram(key='m')
keys = list(hist.keys())
print('{:6}'.format('state'), '|', '{}'.format('count'), '|', '{}'.format('cost'))
for key in keys:
    binary = '{:0=6b}'.format(key)
    count = '{:5}'.format(hist[key])
    print(binary, '|',  count, '|', calc_cost(binary, operators[0][0]))

出力結果例

QAOAパラメータ最適化後の状態に対する
- コストハミルトニアン期待値
- QAOAパラメータ
- 100 shots測定して得られた各状態のカウント数と、今回の制約条件式におけるコスト値(0に近いほど最適解に近い)

今回の問題では状態"100001"と"001100"が共にコスト=-8で最適解であると事前計算で確認可能です。
最適解以外の状態もそれなりにサンプリングされていますが、最適解がより高い確率でサンプリングされることが望ましいです。

Expectation value tf.Tensor([[-22.023201]], shape=(1, 1), dtype=float32)
after params:  [-0.44315836  1.8975906   0.39658815 -3.8769927   4.8495026   1.8638508 ]
state  | count | cost
100001 |     2 | -8.0 # 最適解1
001000 |     8 | -24.0
100000 |     7 | -24.0
010010 |     2 | -16.0
000010 |    15 | -24.0
010100 |     1 | -12.0
011000 |     6 | -50.0
001100 |     6 | -8.0 # 最適解2
000001 |    10 | -24.0
000000 |    13 | -40.0
010000 |    12 | -24.0
000011 |     2 | -50.0
000100 |     8 | -24.0
110101 |     1 | -74.0
100100 |     3 | -16.0
100010 |     2 | -12.0
100011 |     1 | -38.0
101011 |     1 | -72.0

まとめ

TensorFlow-quantumでQAOAをやるメリットがあるかと考えると、個人的には微妙です。

そもそもcirqを今回初めて触ったのですがドキュメントが公式・非公式ともにQISkitなどと比べるとかなり少ないです。
それに変分量子最適化計算をするだけならニューラルネットはscipy.optimizerのような感じで用意された方が取り回しが良いのではという気もします。

本格的に量子深層学習などを実装する場合に真価を発揮するのかもしれませんが、それはハードが追いつかないでしょうしまだ先かなと感じます。

以上、読んで頂きありがとうございました。