Neural QAOA\(^2\): Differentiable Joint Graph Partitioning and Parameter Initialization for Quantum Combinatorial Optimization

Conference: ICML 2026
arXiv: 2605.13072
Code: https://github.com/0SliverBullet/Neural-QAOA-Squared (available)
Area: Quantum Optimization / Differentiable Programming / Graph Partitioning
Keywords: QAOA, divide-and-conquer, differentiable graph partitioning, parameter warm-start, zero-shot generalization

TL;DR

A generative-evaluative neural network (GEN) jointly differentiates the two tasks of "graph partitioning + quantum circuit parameter initialization" in QAOA²: the evaluator learns a high-fidelity quantum performance surrogate, and the generator, guided by its gradients, outputs discrete partitions and parameter initializations. With straight-through estimator (STE) and orthogonal complement head (OCH), the system is end-to-end trainable. Outperforms heuristic baselines on 183 QUBO/Ising/MaxCut instances (21-1000 variables), ranking first on 101 cases.

Background & Motivation

Background: QAOA is the flagship algorithm for QUBO/MaxCut in the NISQ era, but real-world problems often involve thousands of variables, while quantum hardware is limited to hundreds of qubits. The divide-and-conquer paradigm (notably QAOA²) addresses scalability by partitioning large graphs into hardware-fit subgraphs, solving each with QAOA, and merging local solutions via \(\mathbb{Z}_2\) symmetry.

Limitations of Prior Work: Existing D&C frameworks have two decoupled flaws. First, partitioning heuristics (modularity, boundary, KL) are designed for "good graph-theoretic metrics" and are not directly related to final quantum solution quality—the authors report a Pearson correlation of only 0.2859 between modularity and performance ratio on g05_100.1, nearly random. Second, QAOA parameters \((\boldsymbol{\gamma}, \boldsymbol{\beta})\) on subgraphs are randomly initialized, ignoring subgraph topology, resulting in cold-starts—even doubling optimization steps (\(T=40\)) cannot match a topology-aware warm-start (\(T=20\)).

Key Challenge: Both partitioning and parameter initialization are essentially "mapping from graph topology to quantum performance" subtasks, yet are handled separately by heuristics or randomness, with no coordination. End-to-end learning requires solving two engineering challenges: "how to backpropagate through discrete partitions" and "how to enforce hard qubit capacity constraints."

Goal: Construct a differentiable generator that outputs both partitions and initial parameters, with training signals derived from "final quantum performance" rather than intermediate proxy metrics.

Key Insight: Model QAOA² performance prediction as a differentiable surrogate (quantum evaluator), allowing the generator to perform gradient ascent on its output; use STE and greedy capacity discretization (GCD) to "clamp" discrete partitions with hard constraints into the differentiable pipeline; finally, use orthogonal complement head (OCH) to provide a geometric inductive bias for cluster centers, preventing GNN over-smoothing.

Core Idea: Employ a dual-network structure of evaluator and generator to create a differentiable joint strategy for "how to partition and how to initialize," with the evaluator providing quantum-aware gradients, thus truly optimizing D&C for quantum solution outcomes.

Method

Overall Architecture

GEN (Generative Evaluative Network) consists of two parts. First, the Quantum evaluator \(f_\phi(G, \mathbf{S}, \mathbf{P}) \to \hat{\rho}\), a multi-view GNN that encodes the graph \(G\), partition \(\mathbf{S}\), and parameters \(\mathbf{P}\) into a unified latent space to predict the performance ratio \(\rho \in [0.5, 1]\) (formula \(\rho = (\text{Cut} - \text{Neg}) / (\text{OPT} - \text{Neg})\) ensures boundedness). It is first trained to convergence with supervised MSE on a labeled dataset \(\mathcal{D}_{\text{offline}} = \{(G_i, \mathbf{S}_i, \mathbf{P}_i, \rho_i)\}\). Second, the Joint generator \(g_\theta(G) \to (\mathbf{S}, \mathbf{P})\) generates partitions and parameters sequentially via \(P(\mathbf{S}, \mathbf{P} | G) = P(\mathbf{S} | G) P(\mathbf{P} | \mathbf{S}, G)\). With \(f_\phi\) frozen, \(\max_\theta \mathbb{E}_G [f_\phi(G, g_\theta(G))]\) is used for unsupervised gradient ascent.

At inference, a single forward pass \((\mathbf{S}_0, \mathbf{P}_0) = g_\theta(G_{\text{new}})\) provides initial values, followed by test-time adaptation—fine-tuning generator parameters \(\theta\) with a few gradient steps on the specific instance, yielding \(\theta^*\) and outputting \((\mathbf{S}^*, \mathbf{P}^*) = g_{\theta^*}(G_{\text{new}})\).

Key Designs

1. Multi-view Quantum Evaluator \(f_\phi\):

   • Function: High-fidelity prediction of QAOA² performance ratio given partition and initialization, serving as a differentiable surrogate for expensive quantum circuit simulation.
   • Mechanism: Three parallel encoders process heterogeneous inputs—topology encoder for full adjacency \(\mathbf{A}\); partition encoder for subgraph adjacency \(\mathbf{A}_{\text{sub}} = \mathbf{A} \odot (\mathbf{S} \mathbf{S}^T)\) masking inter-partition edges; param encoder broadcasts \(\mathbf{P}\) to node level via \(\mathbf{X}_{\text{param}} = \mathbf{S} \mathbf{P}^T\) and applies sin/cos embedding \(\tilde{\mathbf{X}}_{\text{param}} = [\sin(\mathbf{X}_{\text{param}}), \cos(\mathbf{X}_{\text{param}})]\) to respect \(2\pi\) periodicity. Outputs are globally mean-pooled, concatenated, passed through MLP, and output as \(\hat{\rho} = 0.5(\text{sigmoid}(\text{MLP}(\mathbf{H})) + 1)\) to enforce the theoretical range.
   • Design Motivation: A well-trained differentiable proxy reduces "run quantum simulation = compute gradient" cost from O(simulation) to O(GNN forward pass); multi-view design ensures each input signal has a dedicated encoder, avoiding signal dilution.

2. Orthogonal Complement Head (OCH):

   • Function: Projects GNN node embeddings onto \(k\) cluster centers to obtain soft partition \(\tilde{\mathbf{S}} \in [0,1]^{N \times k}\), while preventing partition probability degeneration due to GNN over-smoothing.
   • Mechanism: Cluster center matrix \(\mathbf{C} \in \mathbb{R}^{k \times h}\) is constrained to two orthogonality conditions \(\mathbf{C} \boldsymbol{g} = \mathbf{0}\) and \(\mathbf{C} \mathbf{C}^T = \mathbf{I}\), where \(\boldsymbol{g} = \text{GMP}(\mathbf{H}_{\text{topology}})\) is the global graph embedding. \(\mathbf{C}\) is dynamically generated via QR decomposition of a random matrix relative to \(\boldsymbol{g}\). Final output is \(\tilde{\mathbf{S}} = \text{softmax}(\mathbf{H}_{\text{topology}} \mathbf{C}^T)\).
   • Design Motivation: Standard GNN + softmax tends to output "almost uniform" partitions due to node embedding similarity (over-smoothing), diluting gradient signals during training. Anchoring cluster centers in the orthogonal complement of the global embedding is equivalent to "subtracting global context," maximizing inter-cluster separability.

3. Greedy Capacity Discretization (GCD) + Straight-Through Estimator (STE):

   • Function: Converts soft \(\tilde{\mathbf{S}}\) into discrete \(\mathbf{S} \in \{0,1\}^{N \times k}\) strictly satisfying qubit capacity constraints (\(\sum_i \mathbf{S}_{ij} \leq \text{max\_nodes}\)), while maintaining gradient flow.
   • Mechanism: GCD greedily assigns nodes to clusters in descending probability order, skipping clusters at capacity, ensuring 100% hard constraint satisfaction. Forward pass uses discrete \(\mathbf{S}\) for precise evaluator scoring; backward pass uses STE \(\nabla_{\tilde{\mathbf{S}}} f \approx \nabla_{\mathbf{S}} f\) to bypass discretization. Parameter generation includes a stop-gradient operator \(\text{sg}(\mathbf{A}_{\text{sub}})\) to prevent parameter optimization from perturbing the partition.
   • Design Motivation: Continuous relaxations like Gumbel-Softmax cannot guarantee hard capacity constraints (qubit count is a physical hardware limit). GCD sacrifices some gradient accuracy for strict feasibility; STE is nearly the only viable differentiable approach for NISQ + discrete decision scenarios.

Loss & Training

Two stages: (1) Evaluator stage minimizes MSE \(\mathbb{E}_{(G, \mathbf{S}, \mathbf{P}, \rho)} [(f_\phi - \rho)^2]\), with data from heuristic partitions, uniformly sampled parameters, and QAOA² simulation ground truth; (2) Generator stage freezes \(f_\phi\) and maximizes \(\mathbb{E}_G [f_\phi(G, g_\theta(G))]\). The generator is trained only at \(p=1\); deeper circuits (\(p=2, 3\)) use Zhou 2020's parameter extension strategy instead of retraining.

Key Experimental Results

Main Results

On 50 held-out test instances (20% holdout from B/BE/W datasets, problem sizes matching training distribution):

Dataset	Random	Modularity	Boundary	KL	Neural QAOA²
B (8 QUBO)	0.8047 (rank 4.75)	0.8351 (2.38)	0.8246 (2.63)	0.8092 (3.75)	0.8417 (1.50, 5/3 wins)
BE (16 QUBO)	0.8626 (4.81)	0.8692 (3.13)	0.8722 (2.31)	0.8672 (3.69)	0.8824 (1.06, 15/1 wins)
W (26 MaxCut)	0.8962 (3.23)	0.9137 (2.23)	0.9114 (2.96)	0.8934 (4.27)	0.9153 (2.23, 8/18 wins)
Overall (50)	0.8708 (3.98)	0.8869 (2.54)	0.8850 (2.70)	0.8716 (4.00)	0.8930 (1.74, 28/22 wins)

On the BE dataset, Neural QAOA² nearly sweeps (15/16), as QUBO typically lacks explicit community structure, rendering modularity-based heuristics ineffective; W (MaxCut) has inherent community structure, so modularity and Neural QAOA² are tied (both rank 2.23).

Ablation Study

93 OOD instances (GKA + L, out-of-distribution, similar scale):

Configuration	GKA (45 QUBO)	L (48 Ising)	Overall (93)
Random	0.8478 (4.16)	0.6984 (4.65)	rank 4.41
Modularity	0.8659 (2.40)	0.7391 (3.06)	rank 2.73
Boundary	0.8601 (2.89)	0.8205 (1.60)	rank 2.24
KL	0.8503 (4.04)	0.7022 (4.27)	rank 4.16
Neural QAOA² (Ours)	0.8762 (1.51, 32/13)	0.8160 (1.42, 28/20)	rank 1.46, 60/33 wins

Zero-shot transfer to OOD topologies (Ising not present in training) also achieves SOTA, indicating GEN learns a general mapping from partition to quantum performance, not dataset-specific features.

Key Findings

• Empirical evidence of low correlation (Pearson 0.2859) between heuristic partitioning and final performance is a key observation supporting the paper's motivation—previously, no one had so clearly quantified this "metric mismatch."
• Even random initialization + \(T=40\) optimization steps cannot outperform topology-aware initialization + \(T=20\) steps, indicating cold-start loss cannot be compensated by more iterations.
• Models trained at \(p=1\) transfer to \(p=2, 3\) and still outperform advanced initialization baselines (TQA/INTERP/FOURIER/QIBPI), showing the learned topological mapping is "parameter schedule agnostic."
• OCH is crucial: removing orthogonal complement constraints in ablation causes GNN outputs to degenerate to nearly uniform distributions, making partition decisions almost random.

Highlights & Insights

• Successfully transforming "heuristic D&C → end-to-end differentiable D&C" is an engineering feat: the authors solve (a) gradient flow for discrete decisions, (b) hard capacity constraints, and (c) GNN over-smoothing—each addressed by a clean component (STE/GCD/OCH), modular and composable.
• The idea of using an evaluator as a differentiable surrogate for gradient signals has precedent in neural architecture search, but its application to quantum combinatorial optimization is novel—it systematically bridges the gap between "expensive oracle evaluation" and "differentiable optimization."
• OCH's dynamic cluster center generation via QR decomposition is clever: traditional clustering treats centers as learnable parameters, prone to collapse with the GNN encoder; anchoring centers in the orthogonal complement of global context provides a fixed geometric anchor.
• Test-time adaptation acknowledges the reality that "training distribution ≠ inference distribution," allowing the model to transition from distributional priors to instance-specific configurations, which is industrially practical.

Limitations & Future Work

• The upper bound of \(\rho\) at 1.0 is relative to the best-known cut, which may not be truly optimal—on large problems, it is provided by classical heuristics, introducing ground-truth bias.
• Both training and evaluation rely on QAOA² simulation, with no validation on real hardware (noise, readout errors, connectivity constraints).
• The generator is trained only at \(p=1\); extension to deeper circuits relies on empirical schedules without theoretical guarantees.
• The hard cap max_nodes=10 does not correspond to the largest current QPUs (hundreds of qubits); scalability to higher qubit counts is untested.
• The training set is 80% B/BE/W; although OOD evaluation is performed on GKA/L, all are from the same benchmark library, with no testing on truly "wild" industrial graphs.

Related Work & Insights

• vs DC-QAOA / original QAOA²: All follow the D&C paradigm, but use heuristic partitioning and random initialization; Neural QAOA² replaces both manual components with end-to-end learnable networks.
• vs INTERP / FOURIER / TQA / QIBPI parameter initialization: These methods only address "how to initialize parameters," with partitioning still heuristic; this work jointly optimizes both, with ablation showing joint training yields much greater gains than parameter-only optimization.
• vs Sampled MuZero / GNN policy in neural combinatorial optimization: All use GNN + RL/differentiable optimization for combinatorial decisions; the novelty here is the evaluator-generator dual network and quantum surrogate, reflecting the higher cost of quantum performance evaluation compared to typical reward signals.
• Insight: Can the evaluator-generator architecture be applied to other "expensive oracle + discrete decision" problems—e.g., chip placement, compiler tuning, hyperparameter search? As long as a high-fidelity differentiable proxy can be learned, the principle is general.

Rating

• Novelty: ⭐⭐⭐⭐ Differentiable D&C is a combinatorial innovation rather than a disruptive new mechanism, but the specific combination of OCH + GCD + STE is a solid contribution
• Experimental Thoroughness: ⭐⭐⭐⭐ 183 instances + 50 IID + 93 OOD + multiple heuristic baselines + different \(p\) depths, systematic comparison
• Writing Quality: ⭐⭐⭐⭐⭐ Motivation diagram (Figure 1) is clear, pipeline description is lucid, key designs have dedicated sections
• Value: ⭐⭐⭐⭐ Directly valuable for NISQ-era quantum combinatorial optimization deployment, code is open-sourced

Related Papers

• [ICML 2026] Learning to Approximate Uniform Facility Location via Graph Neural Networks
• [ICLR 2026] RRNCO: Towards Real-World Routing with Neural Combinatorial Optimization
• [NeurIPS 2025] Probing Neural Combinatorial Optimization Models
• [ICML 2025] Quantum Optimization via Gradient-Based Hamiltonian Descent
• [NeurIPS 2025] Problem-Parameter-Free Decentralized Bilevel Optimization