[논문 리뷰] Optimizing Tensor Network Contraction Using Reinforcement Learning

Eli A. Meirom, Haggai Maron외 2명|arXiv (Cornell University)|2022. 04. 18.

Quantum Computing Algorithms and Architecture인용 수 4

한 줄 요약

이 논문은 양자 회로 시뮬레이션에서 중요한 성능 저하 요인인 텐서 네트워크 수축 순서(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 펄스로 기록하여 Cotengra-Kahypar를 능가했다.
• 8개의 시드와 3시간의 추론 시간을 사용한 결과, RL-TNCO는 최고 기준 솔버 대비 2.5배의 성능 향상을 달성하여 수축 비용을 5.83에서 3.49 × 10^18 펄스로 감소시켰다.
• 대규모 양자 회로에서 시뮬레이션 시간을 2배에서 4배 감소시켜, 회로당 수십만 달러와 GWh 단위의 에너지를 절약하였다.
• 시드 수를 8개 이상 늘릴 경우 성능 향상의 효과가 점점 감소함을 확인하여, 그 이상의 시드 수는 수익 감소의 경향을 보였다.
• 표 8의 초록색 강조 셀을 보면, RL-TNCO는 모든 설정에서 Cotengra-Kahypar를 일관되게 능가했으며, 특히 3시간과 10개의 시드에서 최고 성능를 기록했다.
• 부트스트랩 기반 평가를 통해 중앙값 성능 추정치에 대한 통계적 신뢰도를 확인하였고, 다수의 실행에서 낮은 분산을 보였다.