neural-qaoa-differentiable-optimization

name	neural-qaoa-differentiable-optimization
description	Neural QAOA² methodology: end-to-end differentiable framework for joint graph partitioning and QAOA parameter initialization. Uses generative evaluative network (GEN) with differentiable quantum evaluator for gradient-guided learning. Ranks first on 101/183 instances with zero-shot generalization. Activation: neural QAOA, quantum optimization initialization, graph partitioning QAOA, differentiable quantum, QAOA2.
category	quantum

Neural QAOA²: Differentiable Quantum Optimization

Description

Neural QAOA² is an end-to-end differentiable framework that jointly generates graph partitions and initial parameters for the Quantum Approximate Optimization Algorithm (QAOA). By integrating a generative evaluative network (GEN) with a differentiable quantum evaluator as a high-fidelity performance surrogate, it provides direct gradient guidance, enabling the joint generator to learn the intrinsic mapping from graph topology to high-quality partition and parameter configurations.

arXiv: 2605.13072v1 Authors: Zubin Zheng, Jiahao Wu, Shengcai Liu

Activation Keywords

• neural QAOA
• quantum optimization initialization
• graph partitioning QAOA
• differentiable quantum
• QAOA2
• quantum parameter initialization
• 神经量子优化
• QAOA参数初始化

Core Methodology

Problem Statement

QAOA is promising for combinatorial optimization but constrained by limited qubits. Divide-and-conquer frameworks like QAOA² address scalability by partitioning graphs into subgraphs, but suffer from:

1. Misalignment: Heuristic partitioning metrics don't align with quantum optimization goals
2. Cold starts: Topology-blind parameter initialization leads to optimization inefficiency

Solution: End-to-End Differentiable Framework

Graph Topology → Joint Generator → (Partition, Parameters) → Differentiable Evaluator → Loss → Backprop

Architecture

1. Generative Evaluative Network (GEN)

The GEN has two coupled components:

class GenerativeEvaluatorNetwork(nn.Module):
    def __init__(self, graph_features, partition_dim, param_dim):
        super().__init__()
        self.graph_encoder = GraphEncoder(graph_features)
        self.partition_generator = PartitionGenerator(graph_features, partition_dim)
        self.param_generator = ParameterGenerator(graph_features, param_dim)
        self.quantum_evaluator = DifferentiableQuantumEvaluator()
    
    def forward(self, graph):
        # Encode graph topology
        features = self.graph_encoder(graph)
        
        # Generate partition and parameters jointly
        partition = self.partition_generator(features)
        params = self.param_generator(features)
        
        # Evaluate with differentiable quantum surrogate
        performance = self.quantum_evaluator(graph, partition, params)
        
        return partition, params, performance

2. Differentiable Quantum Evaluator

A high-fidelity surrogate that provides gradient signal:

class DifferentiableQuantumEvaluator(nn.Module):
    def __init__(self, qaoa_depth=2):
        super().__init__()
        self.qaoa_depth = qaoa_depth
    
    def forward(self, graph, partition, params):
        """Differentiable approximation of QAOA performance."""
        # Approximate QAOA expectation value
        # This is differentiable w.r.t. partition and params
        cost = compute_approximate_cost(graph, partition, params)
        return -cost  # Negative because we maximize

3. Joint Generator

Learns the mapping from graph topology to optimal partition and parameters:

class JointGenerator(nn.Module):
    def __init__(self, hidden_dim):
        super().__init__()
        self.gnn = GraphNeuralNetwork(hidden_dim)
        self.partition_head = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, num_partitions)
        )
        self.param_head = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, num_qaoa_params)
        )
    
    def forward(self, graph):
        features = self.gnn(graph)
        partition_logits = self.partition_head(features)
        params = self.param_head(features)
        
        # Use Gumbel-softmax for differentiable partition selection
        partition = gumbel_softmax(partition_logits, tau=0.5, hard=True)
        
        return partition, params

Training Loop

def train_neural_qaoa(model, dataset, epochs=100):
    optimizer = Adam(model.parameters(), lr=1e-3)
    
    for epoch in range(epochs):
        for graph in dataset:
            optimizer.zero_grad()
            
            partition, params, performance = model(graph)
            loss = -performance  # Maximize performance
            
            loss.backward()
            optimizer.step()
            
        if epoch % 10 == 0:
            evaluate(model, validation_set)

Key Results

Performance

Benchmark	Neural QAOA²	Best Heuristic	Improvement
QUBO (101/183)	Rank 1	Heuristic baseline	Varies
Ising	Superior	Standard	Consistent
MaxCut	Superior	Standard	Consistent

Generalization

• Zero-shot generalization: Works on out-of-distribution graph topologies and scales
• Tested range: 21 to 1000 variables
• No retraining needed for new graph types

Application Workflow

Step 1: Prepare Graph Data

import networkx as nx

def prepare_graph(problem_type, n_variables):
    if problem_type == 'maxcut':
        G = nx.erdos_renyi_graph(n_variables, p=0.5)
        weights = {(u, v): 1.0 for u, v in G.edges()}
    elif problem_type == 'ising':
        G = nx.random_regular_graph(3, n_variables)
        weights = {(u, v): np.random.randn() for u, v in G.edges()}
    else:  # QUBO
        G = nx.complete_graph(n_variables)
        weights = {(u, v): np.random.uniform(-1, 1) for u, v in G.edges()}
    
    return G, weights

Step 2: Load Pre-trained Model

model = GenerativeEvaluatorNetwork(
    graph_features=64,
    partition_dim=128,
    param_dim=16
)
model.load_state_dict(torch.load('neural_qaoa2.pth'))
model.eval()

Step 3: Generate Partition and Parameters

partition, params, performance = model(graph)
print(f"Generated partition: {partition}")
print(f"QAOA parameters: {params}")
print(f"Predicted performance: {performance.item()}")

Step 4: Run QAOA

# Use generated partition and parameters to initialize QAOA
qaoa_result = run_qaoa(
    graph=graph,
    subgraphs=partition,
    initial_params=params,
    max_depth=3
)

Design Principles

1. Joint Optimization

Partition and parameter generation are coupled — better partitions enable better parameters and vice versa.

2. Differentiable Surrogate

The quantum evaluator provides direct gradient feedback, eliminating the need for gradient-free optimization.

3. Topology-Aware

The GNN encoder captures graph structure, enabling generalization across different graph types.

4. Scalable

The divide-and-conquer approach allows solving problems with up to 1000 variables on limited qubit hardware.

Error Handling

Poor Convergence

• Increase Gumbel-softmax temperature for smoother gradients
• Use curriculum learning: start with easy graphs, progress to harder ones

Out-of-Distribution Graphs

• The model has demonstrated zero-shot generalization
• If performance degrades, fine-tune on a small sample of new graph type

Parameter Initialization Failure

• Fallback to random initialization with warm start
• Use the generated parameters as a prior for classical optimization

Tools Used

• exec: Run quantum simulation, training scripts
• read: Load graph data, model weights
• write: Save trained models, optimization results

References

• Paper: "Neural QAOA²: Differentiable Joint Graph Partitioning and Parameter Initialization for Quantum Combinatorial Optimization" (arXiv:2605.13072v1)
• QAOA: Quantum Approximate Optimization Algorithm
• Gumbel-softmax: Differentiable discrete sampling