量子GANを使って任意の確率分布を少数ゲートで再現する[実装編]

この記事は量子GANを使って任意の確率分布を少数ゲートで再現する[理論編]の続きで、実装編です。

理論編では
・GANはGenerator部分とDiscriminator部分のニューラルネットで構成され、両者を競合させながら学習している。
・量子GANはGenerator部分を量子ニューラルネットに置き換えたものである。
・量子Generatorのパラメータ$\theta$を最適化することで所望の確率分布を再現する測定結果が得られる。
・パラメータ$\theta$を最適化する過程で古典Discriminatorの学習(パラメータ$\phi$の最適化)が必要
というQGANの性質を詳しく見てきました。

次はこのQGANの実装方法について解説していきます。

実装

ここでの実装はQiskitのQGANチュートリアルをベースにしています。最後に学習過程の可視化とモード崩壊について追記しました。

まず、再現したい確率分布からサンプリングを行います。

import numpy as np
# from qiskit.aqua.utils.dataset_helper import discretize_and_truncate # qiskit 0.25未満
from qiskit_machine_learning.datasets.dataset_helper import discretize_and_truncate # qiskit 0.25以降

N = 1000  # サンプリング数
mu = 1  # 平均1の対数正規分布をロードしていく
sigma = 1  # 標準偏差1の対数正規分布をロードしていく
real_data = np.random.lognormal(mean=mu, sigma=sigma, size=N)

bounds = np.array([0.0, 3.0])  # 抽出区間
num_qubits = [2]  # 量子ビット数

# 分布の離散化
data_trunc, data_grid_trunc, elements_trunc, prob_data = discretize_and_truncate(
    real_data,
    bounds,
    num_qubits,
    return_data_grid_elements=True,
    return_prob=True,
    prob_non_zero=True,
)

次にQGANのクラスを作成します。
ここに再現したい確率分布のサンプリング結果と抽出区間の情報を渡します。
もし多次元分布を再現したい場合は、各次元の抽出区間をbounds = [[min_0,max_0],...,[min_k-1,max_k-1]]で設定し、各次元で使用する量子ビット数をnum_qubits = [#q_0,...,#q_k-1]と設定します。

# from qiskit.aqua.algorithms import QGAN # qiskit 0.25未満
from qiskit_machine_learning.algorithms.distribution_learners.qgan import QGAN # qiskit 0.25以降

batch_size = 100 # バッチサイズ
num_epochs = 10  # エポック数
# Initialize qGAN
qgan = QGAN(real_data, bounds, num_qubits, batch_size, num_epochs)

次に量子Generatorの回路の形を決定します。
チュートリアルでは初期分布として一様分布UniformDistributionを選択していますが、正規分布やランダム生成した回路でも構いません。
また、変分量子回路では使うゲートはRYゲート(パラメータ部分)とCZゲートを選択して繰り返し数を1に設定していますが、パラメータ部分としてRZやRXを使ったり、エンタングル部分でCXゲートを使ってもよいです(選び方によっては収束しない場合もあります)。

from qiskit import BasicAer, ClassicalRegister, QuantumCircuit, execute
from qiskit.circuit.library import TwoLocal
from qiskit.quantum_info import random_statevector
# from qiskit.circuit.library import UniformDistribution, NormalDistribution  # qiskit 0.25未満
from qiskit_finance.circuit.library import UniformDistribution, NormalDistribution # qiskit 0.25以降

# 初期分布作成回路の設定
# 一様分布
init_dist = UniformDistribution(sum(num_qubits))
# 正規分布
# init_dist = NormalDistribution(sum(num_qubits))
# # # ランダム
# init_dist = QuantumCircuit((sum(num_qubits)))
# init_dist.initialize(random_statevector(2**sum(num_qubits)).data, range(sum(num_qubits)))

# 変分量子回路の設定
var_form = TwoLocal(int(np.sum(num_qubits)), "ry", "cz", reps=1)
# 初期分布作成回路と変分量子回路の結合
g_circuit = var_form.compose(init_dist, front=True)

g_circuit.draw("mpl")

初期分布として一様分布回路を選択した場合

> 初期分布として正規分布回路を選択した場合 > 初期分布としてランダム回路を選択した場合

最後に、量子Generatorとして作成した量子回路と、古典のDiscriminatorをQGANクラスに設定します。QiskitライブラリではNumPyで作られたニューラルネットNumPyDiscriminatorとPyTorch版PyTorchDiscriminatorがあります。チュートリアルでは前者が使用されていますが、試した範囲ではPyTorch版の方が挙動が安定しているので今回はこちらを使用します。
これでQGANを実行することができます。出力の相互エントロピーrel_entrが十分に小さくなっていれば確率分布が再現できたとみなせます。

# from qiskit.aqua import QuantumInstance # qiskit 0.25未満
from qiskit.utils import QuantumInstance # qiskit 0.25以降
# from qiskit.aqua.components.neural_networks import NumPyDiscriminator, PyTorchDiscriminator # qiskit 0.25未満
from qiskit_machine_learning.algorithms.distribution_learners import NumPyDiscriminator, PyTorchDiscriminator  # qiskit 0.25以降


init_params = [3.0, 1.0, 0.6, 1.6]  # チュートリアル用に「よい」パラメータを事前に設定
# 量子GeneratorをQGANにセット
qgan.set_generator(generator_circuit=g_circuit, generator_init_params=init_params)

# The parameters have an order issue that following is a temp. workaround
qgan._generator._free_parameters = sorted(g_circuit.parameters, key=lambda p: p.name)

# 古典DiscriminatorをQGANにセット
discriminator = PyTorchDiscriminator(len(num_qubits))
qgan.set_discriminator(discriminator)

# Set quantum instance to run the quantum generator
quantum_instance = QuantumInstance(
    backend=BasicAer.get_backend("statevector_simulator")
)

# QGANを実行
result = qgan.run(quantum_instance)

# 結果の出力
print("Training results:")
print("loss_d :", result["loss_d"])
print("loss_g :", result["loss_d"])
print("rel_entr", result["rel_entr"])

Training results:
loss_d : 0.6925
loss_g : 0.7279
rel_entr : 0.1655

確率分布を再現するための量子Generatorパラメータparams_gを知ってしまえば、次からはこの確率分布を作るときに変分計算(時間を要するパラメータ最適化)を実行する必要がありません。量子回路circuit_gを一回分実行するだけでよいです。得られた分布がチュートリアルと違って見えるのは、チュートリアルでは結果の積算合計を使っているからです。

import matplotlib.pyplot as plt

samples_g, prob_g = qgan.generator.get_output(qgan.quantum_instance, shots=10000)
print("generated prob. :", np.array(prob_g).round(3))
plt.bar(np.array(samples_g).flatten(), prob_g, label="generated")
plt.plot(prob_data, "-o", c="r", label="original")
plt.legend(loc="best")
plt.show()

original prob. : [0.068 0.463 0.337 0.132]
generated prob. : [0.069 0.475 0.306 0.15 ]

DiscriminatorとGeneratorの学習はqgan.run(quantum_instance)で実行できますが、内部モジュールを呼び出すことで、学習の様子を確認することもできます。量子Generator回路の出力の分布がオリジナルデータの分布に近づいていく様子が分かりやすいように、今度は初期パラメータinit_paramsを使わず実行してみます。その分、エポック数は100に設定します。

import random
from tqdm.notebook import tqdm

# Initialize qGAN
qgan = QGAN(real_data, bounds, num_qubits, batch_size, num_epochs, snapshot_dir=None)
qgan.set_discriminator(discriminator)
qgan.set_generator(generator_circuit=g_circuit, generator_init_params=None)

num_epochs = 100
data_all = data_trunc.copy()
for i in tqdm(range(num_epochs)):
    # バッチ数だけ訓練を実行
    index = 0
    np.random.shuffle(data_all)
    while (index + batch_size) <= len(data_all):
        # 1. オリジナルデータのサンプリング
        real_batch = data_all[index : index + batch_size]
        index += batch_size
        # real_probは入力real_batchを調整するための一様分布
        real_prob = np.ones(len(real_batch)) / len(real_batch)

        # 2. 疑似データの生成
        # generated_batchはインデックス
        # generated_probは各インデックスの出現確率
        generated_batch, generated_prob = qgan.generator.get_output(
            quantum_instance, shots=batch_size
        )

        # 3. Discriminatorの訓練
        ret_d = qgan.discriminator.train(
            [real_batch, generated_batch], [real_prob, generated_prob]
        )

        # 4. Generatorの訓練
        qgan.generator.discriminator = qgan.discriminator
        ret_g = qgan.generator.train(quantum_instance, shots=batch_size)

    qgan.d_loss.append(ret_d["loss"])
    qgan.g_loss.append(ret_g["loss"][0])

    if i % int(num_epochs / 10) == 0:
        print("===iter:", i, "===")
        print("---Discriminator---")
        prediction = qgan.discriminator.get_label(
            elements_trunc.reshape(-1, 1)
        ).tolist()
        # 確信度をプリント
        print(
            "confidence that data comes from original for [0, 1, 2, 3]: ",
            np.array(prediction).reshape(-1).round(3),
        )
        print("loss:", ret_d["loss"], "\n")

        print("---Generator---")
        samples_g_out, prob_g_out = qgan.generator.get_output(
            quantum_instance, shots=5000
        )
        # オリジナルの確率分布と生成された確率分布を比較
        print("original prob.      :", np.array(prob_data).round(3))
        print("generated prob. :", np.array(prob_g_out).round(3))
        print("loss:", ret_g["loss"][0].round(8), "\n\n")

        plt.bar(np.array(samples_g_out).flatten(), prob_g_out, label="generated")
        plt.plot(prob_data, "-o", c="r", label="original")
        plt.legend(loc="best")
        plt.show()

===iter: 0 ===
---Discriminator---
confidence that data comes from original for [0, 1, 2, 3]: [0.367 0.545 0.514 0.441]
loss: 0.65936315

---Generator---
original prob. : [0.068 0.463 0.337 0.132]
generated prob. : [0.234 0.255 0.258 0.253]
loss: 0.76813319

:
:
===iter: 90 ===
---Discriminator---
confidence that data comes from original for [0, 1, 2, 3]: [0.388 0.512 0.528 0.459]
loss: 0.7001717

---Generator---
original prob. : [0.068 0.463 0.337 0.132]
generated prob. : [0.035 0.538 0.292 0.135]
loss: 0.68545623

DiscriminatorとGeneratorの目的関数の遷移は次のようになります。

plt.plot(range(num_epochs), qgan.d_loss, label='Discriminator loss')
plt.plot(range(num_epochs), qgan.g_loss, label='Generator loss')
plt.legend(loc='best')
plt.show()

最後にモード崩壊について説明します。モード崩壊とは、Generatorが特定の結果のみを出力してしまう状況で、学習の失敗パターンの１つです。Generatorの学習がDiscriminatorの学習に比べて速すぎる時に発生してしまうことを次のように確かめることができます。極端な例として、Discriminatorの学習を止めて、Generatorだけ学習し続けてみます。

訓練の途中で、Discriminatorが算出した確信度が次のようになっていたとします。

\begin{eqnarray}
D_\phi(x=0)=0.388 \\
D_\phi(x=1)=0.512 \\
D_\phi(x=2)=0.528 \\
D_\phi(x=3)=0.459
\end{eqnarray}

この状況でのGeneratorの目的関数

$$L_G(\phi, \theta) = - \frac{1}{m}\sum_{i=1}^m \left[
\log{D_\phi(\left(G_\theta(z_i)\right)}\right] $$

の最小化ついて考えてみます。上記のようにDiscriminator $D_\phi$の値が定まっていた場合、入力が2であるとき$D_\phi$は最大値0.528をとります。すなわち、Generatorが常に2を出力するように学習が進んで目的関数$L_G(\phi, \theta)$が最小化されてしまいます。この挙動を実際に観測してみます。

num_epochs = 250

# discriminatorのみを学習
for i in tqdm(range(num_epochs)):
    # バッチ数だけ訓練を実行
    index = 0
    np.random.shuffle(data_all)
    while (index + batch_size) <= len(data_all):
        index += batch_size
        ret_g = qgan.generator.train(quantum_instance, shots=batch_size)

    if i % int(num_epochs / 10) == 0:
        print("===iter:", i, "===")
        samples_g_out, prob_g_out = qgan.generator.get_output(
            quantum_instance, shots=5000
        )
        # オリジナルの確率分布と生成された確率分布を比較
        print("original prob.      :", np.array(prob_data).round(3))
        print("generated prob. :", np.array(prob_g_out).round(3))
        print("loss:", ret_g["loss"][0].round(8), "\n\n")

        plt.bar(np.array(samples_g_out).flatten(), prob_g_out, label="generated")
        plt.plot(prob_data, "-o", c="r", label="original")
        plt.legend(loc="best")
        plt.show()      

===iter: 0 ===
original prob. : [0.068 0.463 0.337 0.132]
generated prob. : [0.035 0.538 0.292 0.135]
loss: 0.68545623

:
:
===iter: 245 ===
original prob. : [0.068 0.463 0.337 0.132]
generated prob. : [0. 0.049 0.951 0. ]
loss: 0.65894784

Generatorの出力が2に局在してしまっていることがわかります。QGANの学習を成功させるには、DiscriminatorとGeneratorの学習率を適切に設定する必要があることがわかります。

Version情報

以下の解説で使用したQiskitのVersionは次の通りです。

import qiskit
qiskit.__qiskit_version__

{"qiskit-terra": "0.18.1",
"qiskit-aer": "0.8.2",
"qiskit-ignis": "0.6.0",
"qiskit-ibmq-provider": "0.16.0",
"qiskit-aqua": "0.9.4",
"qiskit": "0.29.0",
"qiskit-nature": "0.1.5",
"qiskit-finance": "0.2.0",
"qiskit-optimization": "0.2.1",
"qiskit-machine-learning": "0.2.0"}