🤖 machine learning

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

This paper introduces a learning-to-rank framework that utilizes gradient-boosted models trained on GPU performance data to efficiently select optimal tensor network contraction plans for quantum circuit simulation, demonstrating improved decision quality over traditional baselines while maintaining robustness across different GPU architectures.

Original authors: Alfred M. Pastor, Maribel Castillo, Jose M. Badia

Published 2026-08-07

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Alfred M. Pastor, Maribel Castillo, Jose M. Badia

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve a massive, impossible puzzle, but instead of cardboard pieces, you are juggling invisible, multi-dimensional shapes called "tensors." This is the world of quantum computing, where scientists try to simulate how quantum computers think. The problem is, these simulations are incredibly expensive to run on regular computers; the more complex the puzzle, the more memory and time it eats up, often growing so fast that it becomes impossible to finish. To make this manageable, researchers use a clever trick called "tensor networks," which breaks the giant puzzle into smaller, connected chunks. However, there's a catch: there are billions of ways to put these chunks back together. Some ways are like taking a scenic, winding road that takes forever, while others are like a highway that zooms you to the finish line. Finding the fastest route is a nightmare for computers to solve on their own. This is where the magic of "Learning to Rank" comes in—a technique borrowed from how search engines decide which websites to show you first. Instead of trying to calculate the exact time for every single route (which would take too long), the goal is to quickly guess which few routes are the best, so you only have to test those.

This paper is about teaching a computer to become a super-smart tour guide for these quantum puzzle routes, specifically for high-speed graphics cards (GPUs) that are used to do the heavy lifting. The researchers, Alfred M. Pastor, Maribel Castillo, and Jose M. Badia, realized that while two routes might look similar on paper, they can perform wildly differently when actually run on a GPU, depending on how the card handles memory and parallel work. So, they built a system that looks at the "shape" of a contraction plan (the route) and uses a machine learning model to rank them from "best guess" to "worst guess" without ever actually running the simulation first.

Here is how they did it: They took 225 different quantum circuit puzzles and generated seven different possible routes for each one. They ran all these routes on an NVIDIA RTX A6000 graphics card to see which ones were actually the fastest. Then, they fed the "structural features" of these routes—things like how much work each step requires, how balanced the shapes are, and how much data needs to be moved—into a smart algorithm called XGBoost. They taught this algorithm to learn which features predicted the fastest times. They tested two different ways of teaching it: one that focused on getting the very top spot right (like a listwise approach) and another that focused on comparing pairs of routes (like a pairwise approach).

The results were quite promising. When tested on new puzzles that looked like the ones they learned from, the "listwise" model was a star. It managed to put the actual fastest route in its top three guesses for 96% of the circuits. Even better, it was often the only one that could correctly pick the single fastest route as its number one choice 60% of the time. This is a huge improvement over just guessing randomly or using older, standard methods that didn't learn from data. The model learned that the "shape" of the route matters just as much as the raw math, capturing things like how well the route fits the GPU's memory habits.

However, the paper also sounds a cautious note about how much this "tour guide" can travel. When they tested the model on a completely different family of puzzles (Quantum Fourier Transform circuits) that it had never seen before, its performance dropped. It still found good routes 62.9% of the time in its top three, but it wasn't as perfect as before. This suggests that while the model is smart, it still relies on seeing similar types of puzzles to work its best.

They also asked a fascinating question: "If we train this guide on one type of graphics card (the Ampere architecture), will it still work if we switch to a different type (the Volta architecture) without retraining?" The answer was a cautious "yes, but with some bumps." The two different graphics cards agreed on the absolute fastest route 84% of the time. When the model trained on the first card was tested on the second, it still found a good route in its top three 92% of the time. This suggests that the structural features of the routes are somewhat universal, but the specific hardware does change the outcome enough that the model isn't perfectly portable yet.

In short, the authors have shown that you can use machine learning to skip the expensive trial-and-error phase of finding the best quantum simulation routes. By looking at the blueprint of the route, the model can suggest a short list of the best candidates, saving a massive amount of time. While it's not a magic wand that works perfectly on every single puzzle or every single computer chip, it is a practical tool that significantly reduces the search effort, making it easier to simulate larger and more complex quantum circuits on today's hardware.

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

Problem Statement
Classical simulation is vital for validating quantum algorithms, but its cost scales exponentially with circuit size. Tensor-network methods offer a solution by exploiting circuit structure to reduce memory and computational requirements. However, the efficiency of a tensor-network simulation is critically dependent on the contraction plan—the specific sequence in which tensors are contracted. Finding an optimal contraction order is NP-hard; thus, practical tools rely on heuristics (e.g., MinFill, graph partitioning) to generate candidate plans.

A significant gap exists between theoretical cost estimates and actual execution performance on GPUs. While plans may have similar theoretical complexity (e.g., floating-point operation counts), their runtime on GPU hardware varies drastically due to factors such as parallelism, memory traffic, contraction geometry, and kernel granularity. Measuring all candidate plans to find the fastest one negates the efficiency gains of the selection process. The authors formulate the problem as selecting the plan expected to minimize execution time (or a short list of near-optimal candidates) without executing them first.

Methodology
The authors propose a Learning-to-Rank (LTR) framework to select contraction plans based on structural features rather than runtime profiling.

Feature Engineering:
- The input representation is derived entirely from the contraction plan's sequence of pairwise contractions, requiring no preliminary execution.
- The authors define a cost model based on base-2 logarithmic coordinates for local contraction steps, identifying four primitives: computational burden ( $c_i$ ), output size/parallelism ( $p_i$ ), reduction dimension ( $k_i$ ), and contraction geometry imbalance ( $d_i$ ).
- These local primitives are aggregated into a fixed-size vector of 18 features (later reduced to 14-15 via selection) covering five groups: complexity/critical load, sequence/granularity, parallelism dynamics, reduction/intensity, and geometry/memory access.
Learning Framework:
- The task is treated as a grouped LTR problem, where candidates for a single circuit form a ranking group.
- The target variable is a relevance label derived from the measured speedup of a candidate plan relative to a MinFill baseline on a specific GPU, transformed via a logarithm to create a graded relevance score.
- Models: The authors train gradient-boosted rankers using XGBoost. Two objectives are compared:
  - Pairwise: Minimizes ranking inversions (RankNet-style).
  - Listwise: Directly optimizes Normalized Discounted Cumulative Gain (NDCG), specifically focusing on the top positions of the ranking.
- Baselines: The models are compared against random ranking, a pure MinFill strategy, and hybrid strategies that force MinFill to the top position.
Experimental Setup:
- Dataset: 225 circuit groups with seven candidate plans each, generated by diverse heuristics (MinFill, FlowCutter, community detection, etc.).
- Hardware: Models are trained on an NVIDIA RTX A6000 (Ampere) and evaluated on both the A6000 and an NVIDIA Tesla V100 (Volta) to study backend shift.
- Validation: A strict partitioning strategy prevents data leakage, using a locked in-distribution (ID) test set and an out-of-distribution (OOD) set consisting entirely of a held-out QFT-based circuit family.

Key Contributions

Formulation: The paper formulates GPU contraction-plan selection as a grouped Learning-to-Rank problem, aiming to identify a Top-1 or Top-k candidate before execution.
Representation: It defines a fixed-size structural representation of contraction plans that captures GPU-relevant aspects (work, parallelism, geometry) without requiring timing measurements at inference time.
Protocol: The authors develop a group-aware training and evaluation protocol for XGBoost rankers, utilizing both ranking accuracy and decision-oriented regret metrics across ID and OOD splits.
Backend Analysis: The study examines backend dependence through two analyses: the stability of empirical ground-truth rankings across Ampere and Volta architectures, and the zero-shot transfer of Ampere-trained models to Volta.

Results

In-Distribution Performance: On the locked ID test set, the NDCG-oriented model (opt_ndcg) achieved a Top-3 accuracy of 96%, meaning the fastest plan was within the top three recommendations for 96% of circuits. It significantly outperformed random and MinFill-based baselines. The Top-1 accuracy was 60%.
Decision Quality: The NDCG model demonstrated lower decision regret (the loss in speedup when selecting a suboptimal plan) compared to the pairwise model, particularly in the Top-3 selection.
Circuit-Family Shift (OOD): When evaluated on the held-out QFT-based family (completely unseen during training), performance degraded but remained useful. The Top-3 rate for the NDCG model dropped to 62.9%, indicating that while the model retains predictive signal for unseen structures, generalization is sensitive to the specific circuit family.
Backend Shift (Zero-Shot):
- Ground Truth Stability: Empirical rankings on Ampere and Volta agreed on the same fastest plan in 84% of circuits and shared 89.6% of the Top-3 plans on average.
- Model Transfer: When Ampere-trained models were evaluated on Volta labels without retraining, the NDCG model maintained a Top-3 rate of 92%. However, regret increased, indicating that while the shortlist remains useful, the relative ordering of the very best candidates is somewhat backend-dependent.

Significance and Claims
The paper claims that structural information derived from contraction plans is sufficient to support useful ranking decisions for GPU-accelerated quantum simulation. The results support Learning-to-Rank as a practical method to reduce the search space of contraction plans, thereby saving the computational cost associated with measuring multiple candidates.

The authors maintain a modest stance regarding hardware independence. They conclude that while the learned orderings are substantially stable across the two tested NVIDIA architectures (Ampere and Volta), the performance is partly backend-dependent. The observed stability does not establish general hardware portability, as the study is limited to two GPUs with closely related software stacks and a specific set of candidate heuristics. The method is presented as a low-cost selection stage to be placed between plan generation and execution, rather than a generator of new contraction orders. Future work is proposed to include larger candidate sets, stricter domain-shift experiments, and potentially conditioning rankers on specific hardware descriptors.

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

More like this