[Paper Review] Optimizing Tensor Network Contraction Using Reinforcement Learning

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

Quantum Computing Algorithms and Architecture4 citations

TL;DR

This paper proposes a reinforcement learning (RL) approach using graph neural networks (GNNs) to optimize tensor network contraction ordering (TNCO), a critical bottleneck in quantum circuit simulation. By formulating TNCO as a Markov Decision Process and training an RL agent to minimize contraction cost, the method achieves 2x–4x speedups over state-of-the-art baselines, significantly reducing computational time and energy use in large-scale quantum simulations.

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.

Motivation & Objective

• To address the high computational cost of tensor network contraction in quantum circuit simulation, which depends critically on contraction order.
• To formulate the tensor network contraction ordering (TNCO) problem as a reinforcement learning (RL) Markov Decision Process (MDP).
• To design a GNN-based RL agent that learns to select contraction sequences minimizing total computational cost.
• To overcome challenges including massive search space, heavy-tailed rewards, and difficult credit assignment in long episodes.
• To outperform existing state-of-the-art solvers, especially on large-scale quantum circuits like Sycamore M20.

Proposed method

• The TNCO problem is modeled as an MDP where the state is a tensor network graph, actions are edge contractions, and the reward is the negative contraction cost.
• A graph neural network (GNN) serves as the policy network, outputting a probability distribution over edges to contract at each step.
• The agent is trained using Proximal Policy Optimization (PPO) to minimize cumulative contraction cost over entire contraction sequences.
• To handle heavy-tailed rewards and long-horizon credit assignment, the method employs careful curriculum learning and experience replay with bootstrapped evaluation.
• The approach integrates with existing solvers by using their outputs as initialization or supervision signals during training.
• Multiple random seeds are used per configuration, and the best result across seeds is reported to account for stochasticity in RL and baseline solvers.

Experimental results

Research questions

• RQ1 Can a learning-based RL approach outperform non-learning-based state-of-the-art solvers in tensor network contraction ordering?
• RQ2 How does the proposed RL-GNN method handle the massive search space and heavy-tailed reward distribution inherent in TNCO?
• RQ3 What is the marginal benefit of increasing inference time and number of seeds for the RL-based solver compared to deterministic baselines?
• RQ4 To what extent can RL-based contraction ordering reduce the computational cost and energy consumption in large-scale quantum circuit simulations?
• RQ5 Can the RL agent generalize across different tensor network architectures, including those from real-world quantum circuits?

Key findings

• The proposed RL-TNCO method achieves a median contraction cost of 3.49 × 10^18 flops on the Sycamore M20 circuit with 3 hours of inference and 10 seeds, outperforming Cotengra-Kahypar.
• With 8 seeds and 3 hours of inference, RL-TNCO achieves a 2.5x speedup over the best baseline, reducing contraction cost from 5.83 to 3.49 × 10^18 flops.
• The method reduces simulation time by 2x–4x on large-scale quantum circuits, saving hundreds of thousands of dollars and GWhs of energy per circuit.
• The performance gain from increasing the number of seeds diminishes after 8 seeds, indicating diminishing returns beyond that point.
• Green-highlighted cells in Table 8 show that RL-TNCO consistently outperforms Cotengra-Kahypar across all configurations, with the best results achieved at 3 hours and 10 seeds.
• The bootstrap-based evaluation confirms statistical confidence in the median performance estimates, with low variance across multiple runs.