[论文解读] 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

一句话总结

本文提出一种基于图神经网络（GNN）的强化学习（RL）方法，用于优化张量网络收缩顺序（TNCO），这是量子线路模拟中的关键性能瓶颈。通过将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）作为策略网络，输出每一步需收缩的边的概率分布。
• 使用近端策略优化（PPO）训练智能体，以最小化整个收缩序列的累积收缩成本。
• 为处理重尾奖励分布和长时程信用分配问题，采用细致的课程学习策略，并结合基于Bootstrap的评估机制进行经验回放。
• 通过在训练过程中使用现有求解器的输出作为初始化或监督信号，将该方法与现有求解器集成。
• 每个配置使用多个随机种子，报告各种子中的最佳结果，以考虑强化学习和基线求解器中的随机性。

实验结果

研究问题

• RQ1 基于学习的RL方法能否在张量网络收缩顺序任务中超越非学习型的最先进求解器？
• RQ2 所提出的RL-GNN方法如何处理TNCO中固有的巨大搜索空间和重尾奖励分布问题？
• RQ3 与确定性基线相比，增加推理时间与种子数量对基于RL的求解器带来的边际收益如何？
• RQ4 基于RL的收缩顺序优化在大规模量子线路模拟中，能在多大程度上降低计算成本和能耗？
• RQ5 RL智能体能否在不同张量网络架构上泛化，包括来自真实世界量子线路的架构？

主要发现

• 在Sycamore M20电路上，所提出的RL-TNCO方法在3小时推理时间与10个种子条件下，实现了3.49 × 10^18 次浮点运算的中位收缩成本，优于Cotengra-Kahypar。
• 在8个种子与3小时推理时间下，RL-TNCO相比最佳基线实现了2.5倍的速度提升，将收缩成本从5.83 × 10^18 降低至3.49 × 10^18 次浮点运算。
• 该方法在大规模量子线路中将模拟时间减少了2倍至4倍，每条电路可节省数十万美元及数百万千瓦时的能源。
• 当种子数量超过8个后，性能提升趋于平缓，表明在此之后收益递减。
• 表8中绿色高亮的单元格显示，RL-TNCO在所有配置下均持续优于Cotengra-Kahypar，最佳结果出现在3小时与10个种子条件下。
• 基于Bootstrap的评估验证了中位性能估计的统计置信度，多轮运行中表现出较低方差。