Learning to Rank Tensor Network Contraction Plans for GPU-Accelerated Quantum Circuit Simulation

As quantum circuits grow in size, classical simulation via tensor-network contraction becomes computationally expensive. The efficiency of this process is highly sensitive to the chosen contraction plan—the specific sequence of pairwise tensor operations. Because finding an optimal plan is NP-hard and performance on GPUs depends on complex factors like memory traffic and parallelism rather than just floating-point counts, the authors investigate whether machine learning can predict the most efficient contraction plan for a given circuit before execution.

The authors frame the problem as a grouped Learning to Rank (LTR) task. They represent each contraction plan using structural features derived from its sequence of pairwise contractions, capturing aspects like computational load, reduction structure, and contraction geometry. They train gradient-boosted rankers (using XGBoost) with both pairwise and listwise (NDCG-oriented) objectives. The models are trained on measured GPU execution times and evaluated on their ability to rank candidate plans for both in-distribution circuits and a held-out circuit family (QFT-based circuits).

The study demonstrates that structural features are sufficient to predict relative plan performance. The listwise NDCG-oriented model outperformed other approaches, successfully including the fastest contraction plan in its top-three recommendations for 96% of the in-distribution test cases. When tested on the out-of-distribution QFT circuit family, the model's top-3 accuracy dropped to 62.9%, indicating that while the model generalizes to some extent, it remains sensitive to circuit structure shifts. Furthermore, the authors found that empirical rankings are relatively stable across different GPU architectures (NVIDIA RTX A6000 and Tesla V100), allowing models trained on one architecture to provide useful, though not perfect, guidance on another.