[論文レビュー] Optimizing Tensor Network Contraction Using Reinforcement Learning

Eli A. Meirom, Haggai Maron+2 more|arXiv (Cornell University)|Apr 18, 2022

Quantum Computing Algorithms and Architecture被引用数 4

ひとこと要約

本論文は、量子回路シミュレーションにおける重要なボトル neck であるテンソルネットワーク縮約順序（TNCO）を最適化するため、グラフニューラルネットワーク（GNN）を用いた強化学習（RL）アプローチを提案する。TNCOをマークフ・決定過程（MDP）として定式化し、縮約コストを最小化するようにRLエージェントを訓練することで、最先端のベースラインと比較して2倍～4倍の高速化を達成し、大規模な量子シミュレーションにおける計算時間とエネルギー消費量を顕著に削減する。

ABSTRACT

Quantum Computing (QC) stands to revolutionize computing, but is currently still limited. To develop and test quantum algorithms today, quantum circuits are often simulated on classical computers. Simulating a complex quantum circuit requires computing the contraction of a large network of tensors. The order (path) of contraction can have a drastic effect on the computing cost, but finding an efficient order is a challenging combinatorial optimization problem. We propose a Reinforcement Learning (RL) approach combined with Graph Neural Networks (GNN) to address the contraction ordering problem. The problem is extremely challenging due to the huge search space, the heavy-tailed reward distribution, and the challenging credit assignment. We show how a carefully implemented RL-agent that uses a GNN as the basic policy construct can address these challenges and obtain significant improvements over state-of-the-art techniques in three varieties of circuits, including the largest scale networks used in contemporary QC.

研究の動機と目的

• テンソルネットワーク縮約の高い計算コストが、縮約順序に強く依存することに起因する課題を解決すること。
• テンソルネットワーク縮約順序（TNCO）問題を強化学習（RL）のマークフ・決定過程（MDP）として定式化すること。
• 総計算コストを最小化する縮約シーケンスを選択できるGNNベースのRLエージェントを設計すること。
• 膨大な探索空間、重尾型の報酬分布、および長時間にわたるエピソードにおける困難な責任割り当てといった課題に直面する問題を克服すること。
• 特にSycamore M20のような大規模な量子回路において、既存の最先端ソルバーを上回ること。

提案手法

• TNCO問題は、状態がテンソルネットワークグラフ、行動がエッジの縮約、報酬が縮約コストの負の値であるMDPとしてモデル化される。
• ポリシー関数としてグラフニューラルネットワーク（GNN）が用いられ、各ステップで縮約するエッジの確率分布を出力する。
• エージェントは、全縮約シーケンスにおける累積的な縮約コストを最小化するように、Proximal Policy Optimization（PPO）を用いて訓練される。
• 重尾型報酬と長時間にわたる報酬割り当ての問題に対処するため、慎重なカリキュラム学習とブートストラップ評価を用いた経験再生を採用する。
• 訓練中に既存のソルバーの出力を初期化または監督信号として統合することで、それらのソルバーと連携可能である。
• 各設定に対して複数のランダムシードを用い、シード間の最良の結果を報告することで、RLおよびベースラインソルバーの確率的性質を反映する。

実験結果

リサーチクエスチョン

• RQ1 学習ベースのRLアプローチは、非学習ベースの最先端ソルバーを上回ることができるか？
• RQ2 提案されたRL-GNN手法は、TNCOに内在する巨大な探索空間と重尾型報酬分布をどのように処理するか？
• RQ3 RLベースのソルバーにおいて、推論時間とシード数を増加させた際の限界的利益は、決定的ベースラインと比較してどの程度か？
• RQ4 RLベースの縮約順序決定は、大規模な量子回路シミュレーションにおける計算コストとエネルギー消費をどの程度削減できるか？
• RQ5 RLエージェントは、実世界の量子回路から得られるような異なるテンソルネットワークアーキテクチャにも一般化可能か？

主な発見

• Sycamore M20回路において、3時間の推論時間と10個のシードを用いた場合、提案されたRL-TNCO手法は中央値の縮約コストが3.49 × 10^18 flopsに達し、Cotengra-Kahyparを上回った。
• 8個のシードと3時間の推論時間で、RL-TNCOは最良のベースラインを2.5倍のスピードアップで上回り、縮約コストを5.83から3.49 × 10^18 flopsに削減した。
• 大規模な量子回路において、シミュレーション時間を2倍～4倍短縮し、1回の回路あたり数十万ドルおよびGWh単位のエネルギーを節約した。
• シード数を8個を超えて増加させても性能向上の割合が減少し、その点以降はリターンが逓減することが示された。
• 表8の緑色でハイライトされたセルは、RL-TNCOが全設定でCotengra-Kahyparを一貫して上回っていることを示しており、特に3時間と10個のシードで最良の結果が得られた。
• ブートストラップ評価により、中央値の性能推定値に対する統計的信頼性が確認され、複数回の実行においても低い分散を示した。