Technical Summary: Rethinking Efficiency in Neural Combinatorial Optimization: Batched Preference Optimization with Mamba

1. Problem Statement

Neural Combinatorial Optimization (NCO) has evolved from early Pointer Networks to Transformer-based architectures (e.g., AM, POMO) and Generative Flow Networks. While these methods have improved solution quality, they face significant efficiency bottlenecks, particularly regarding scalability and training convergence speed.

The paper identifies two primary sources of inefficiency in existing NCO works:

1. Learning Paradigm: Traditional online Reinforcement Learning (RL) tightly couples trajectory generation (rollouts) with parameter updates. This "sample-then-update" loop under-utilizes GPU parallelism and forces every optimization step to depend on fresh, autoregressive generation, leading to slow convergence and high hardware costs.
2. Neural Architecture: State-of-the-art NCO models rely heavily on the Transformer architecture. Despite optimizations like FlashAttention, the underlying token-token interaction remains quadratic ($O(N^2)$) in sequence length. This causes memory explosion and latency issues on long routing instances (e.g., TSP with thousands of nodes), limiting the practical deployment of NCO on large-scale problems.

The core question addressed is: How can we improve NCO efficiency without sacrificing solution quality?

2. Methodology: The ECO Framework

The authors propose ECO, an efficiency-oriented framework that decouples trajectory generation from gradient updates and replaces the Transformer backbone with a Mamba-based architecture.

A. Batched Preference Optimization Paradigm

ECO introduces a two-stage training workflow that shifts from strict online RL to a "batched or offline-in-update" paradigm:

1. Phase I: Supervised Warm-up (SFT): The policy is initialized via behavior cloning on high-quality pre-computed solutions (generated by LKH-3 or heuristics). This provides a basic understanding of routing constraints and local optimality.
2. Phase II: Iterative Direct Preference Optimization (DPO): Instead of training a separate value function (Critic) as in PPO, ECO uses DPO to optimize the policy directly from pairwise preferences.
   • Workflow: The current policy generates a batch of candidate solutions. Preference pairs (Winner/Loser) are constructed based on tour length.
   • Reference Model: A reference model ($\pi_{ref}$) is updated periodically (every $T$ steps) to anchor the KL divergence, stabilizing the iterative improvement.
   • LS-Augmented Bootstrapping: To prevent vanishing gradients caused by narrow reward margins as the policy improves, a lightweight local search (LS) operator (e.g., 2-opt, 3-opt) is applied to model-generated solutions during training. The refined solution becomes the "Winner," and the raw output becomes the "Loser." Crucially, local search is never used at inference time; it serves only to construct stronger training signals that are distilled into the neural policy.

B. Mixed Mamba-based Architecture

To address the quadratic memory scaling of Transformers, ECO employs a tailored Mamba (State Space Model) architecture:

1. Recurrent Mamba Encoder: Uses stacked Mamba layers (S6 setting) to process node sequences. Unlike Transformers, Mamba uses linear-memory sequence modeling. To handle the non-permutation-equivariant nature of Mamba, inputs are canonicalized (e.g., sorted by Morton/Z-order code) before encoding.
2. S4 Decoder: Uses S4-style layers for parallel teacher-forced computation during training and efficient recurrent state updates during inference.
3. Autoregressive Scoring: The framework retains the standard autoregressive candidate-scoring step (pointer-style compatibility scoring). The authors explicitly note that while the encoder and recurrent state updates become linear, the action scoring remains quadratic in the worst case ($O(N^2)$) because it requires comparing the current state against all feasible candidates. However, the architecture significantly reduces encoder memory pressure and recurrent state costs.

3. Key Contributions

The paper outlines three primary contributions:

1. First Mamba-based NCO Solver: The authors introduce the first NCO solver based on Mamba, providing a complexity decomposition showing that Mamba reduces encoder memory scaling and recurrent decoding-state costs, though naive autoregressive candidate scoring remains quadratic.
2. Efficiency-Oriented Framework: A novel integration of a tailored Mamba architecture, a two-stage batched preference optimization pipeline (SFT + Iterative DPO), and a local-search-guided bootstrapping strategy to stabilize training and amplify preference margins.
3. Comprehensive Empirical Analysis: Extensive case studies on TSP and CVRP demonstrating that ECO achieves best-or-competitive neural performance with substantial efficiency gains, particularly on large instances. The analysis includes complexity decomposition, permutation sensitivity tests, and ablation studies on each design component.

4. Experimental Results

Experiments were conducted on TSP ($N \in \{200, 500, 1000, 5000\}$) and CVRP ($N \in \{100, 200, 500, 1000\}$) using a single NVIDIA A800 GPU.

• Solution Quality: ECO achieves the best or competitive results among neural baselines (AM, POMO, LEHD, CNF, GFlowNet-HBG) across most scales.
  • On TSP5000, ECO achieves a gap of 5.49%, whereas Transformer-based baselines (AM, POMO) suffer severe degradation with gaps exceeding 70-89%.
  • On CVRP1000, ECO achieves a gap of 2.87%, outperforming or matching other neural methods.
• Scalability and Memory:
  • Memory: The Transformer baseline triggers Out-Of-Memory (OOM) errors at $N \approx 7000$ on an 80GB GPU due to quadratic attention memory. ECO maintains near-linear memory growth, staying within limits even at $N=10,000$ (approx. 40GB).
  • Throughput: ECO sustains significantly higher inference throughput on large sequences compared to Transformers, as it can utilize larger batch sizes without hitting memory limits.
• Training Efficiency:
  • ECO converges much faster than standard online baselines.
  • Comparing DPO to PPO with the same Mamba backbone, ECO achieves comparable quality (4.84% vs 4.88% gap on TSP1000) but reduces training time from 34.5 hours to 10.2 hours.
• Component Analysis:
  • LS Bootstrapping: Removing LS guidance increases the TSP1000 gap from 4.84% to 9.55%, confirming that LS provides a stronger training signal that is distilled into the policy.
  • SFT Warm-up: Training without SFT leads to instability and divergence, confirming the need for a competent initial policy.
  • Architecture: Replacing Mamba with Transformer under the same DPO objective increases the gap (5.47% vs 4.84%) and training time, proving the architectural benefit is independent of the optimization objective.

5. Significance and Claims

The paper claims that ECO represents a significant step toward scalable, practical neural routing solvers by addressing the "efficiency" gap that has been largely overlooked in favor of solution quality.

• Modest Claims on Linearity: The authors are explicit that ECO does not make the entire constructive inference linear in $N$. The benefit is specific: it reduces encoder memory scaling and recurrent state costs, and improves hardware utilization via scan kernels and batched updates. The final action scoring step remains quadratic.
• Training vs. Inference: A key finding is that local search is strictly a training-time mechanism. The performance gains of ECO over ablated versions are attributed to improved network parameters learned from LS-guided signals, not to post-processing at inference.
• Hardware Awareness: The work demonstrates that combining hardware-aware architectures (Mamba) with efficient training paradigms (Batched DPO) allows NCO models to scale to thousands of nodes on consumer-grade or single-GPU hardware, where traditional attention-based models fail due to memory constraints.

In conclusion, ECO provides a strong direction for scalable neural routing by unifying a linear-memory backbone with a stable, batched preference-learning pipeline, achieving competitive solution quality with markedly better memory and throughput profiles.