Probing Neural Combinatorial Optimization Models

AI summary

Overview

Research area: Interpretability of neural combinatorial optimization (NCO) — specifically, applying linear probing techniques from computer vision and NLP to understand what transformer-based NCO solvers learn internally.

Technical level: Intermediate. Readers should be comfortable with attention-based neural architectures, linear probing methodology, and routing problem formulations (TSP, CVRP), but the paper's framing is accessible to graduate students in ML or operations research.

Scope: The paper designs probing tasks for NCO models, introduces a new analysis tool called Coefficient Significance Probing (CS-Probing), and uses it to explain inductive biases, generalization behavior, and key embedding dimensions in three representative models (AM, POMO, and LEHD).

What This Paper Is About

Neural combinatorial optimization (NCO) models such as AM, POMO, and LEHD now match or exceed specialized heuristics like Concorde, ACO, LKH3, and HGS, yet nobody knows what knowledge these models actually encode in their hidden representations. The authors take a first systematic step toward interpreting NCO models by probing their learned embeddings with auxiliary linear prediction tasks designed specifically for combinatorial optimization, augmented by a new tool (CS-Probing) that examines both the magnitude and statistical significance of probing coefficients. The goal is to answer two questions: what decision-related knowledge do NCO models acquire, and how do they learn and use it.

Key Contributions

1. Systematic probing tasks for NCO. Because combinatorial optimization problems lack the natural subtasks available in NLP, the authors design four probing tasks: two for TSP (perceiving Euclidean distance, and avoiding myopic nearest-node decisions) and two for CVRP (encoding the linear additive relationship between node demands, and encoding whether two nodes belong to the same route in the optimal solution).

2. A new probing tool, CS-Probing. Coefficient Significance Probing analyzes both the absolute magnitude and the statistical significance of each dimension's coefficient in a linear probe, enabling fine-grained, per-dimension analysis of what individual embedding neurons encode.

3. Direct evidence about inductive biases and generalization. CS-Probing reveals that AM, POMO, and LEHD impose distinct inductive biases on their representations, and shows that models that generalize well (LEHD) consistently reuse the same embedding dimensions across tasks and problem scales, while models whose knowledge becomes disorganized across dimensions during generalization suffer performance degradation.

4. Practical payoff from probing insights. Using the identified key dimensions, the authors show that LEHD retains near-equivalent performance using only two of its 128 output dimensions, and that adding a sparsity-promoting regularization term — a minor code modification — improves generalization at larger problem scales.

Main Findings

• Initial embeddings carry no linearly decodable distance information. For AM-Init without interaction terms, the probe regression R² was -0.0003 ± 0.00000 and Task 2 AUC was 0.49 ± 0.00. The paper reports that an R² of zero for initial embeddings holds across all NCO models, because raw 2D coordinates are mapped through a linear projection into a 128-dimensional shared space and Euclidean distance is nonlinear.

• After passing through the model, distance becomes linearly decodable. R² rose to 0.2529 ± 0.00048 for AM-Enc-l3, 0.1981 ± 0.00001 for POMO-Enc-l6, and 0.9418 ± 0.00031 for LEHD-Dec-l6 (20-node instances, without interaction terms). With interaction terms, all three models' post-encoder or post-decoder embeddings approached R² of 1 (AM-Enc-l3: 0.9282 ± 0.00035; POMO-Enc-l6: 0.7917 ± 0.00000; LEHD-Dec-l6: 0.9415 ± 0.00027).

• All three models learn to avoid myopic decisions. On Probing Task 2, AM, POMO, and LEHD all achieved AUC scores exceeding 0.8, and this ability was consistently stronger in the final layer than the first layer for all three models. A probe trained on raw path distance features alone yielded AUC close to 0.5, confirming the task is not trivial.

• The two probing tasks are fundamentally different. Initial embeddings with interaction terms achieved R² above 0.7 on distance perception yet AUC near 0.5 on myopia avoidance, showing that the second task is not simply a restatement of Euclidean distance.

• NCO models shift from spatial perception to strategic reasoning with depth. Initial embeddings showed weak Euclidean distance perception; after just one attention layer, all models achieved strong distance awareness, which slightly weakened with depth, while deeper layers developed high-level abilities such as avoiding myopic node choices. LEHD's final layers showed a slight decline in this ability, which the authors attribute to possibly more complex strategies.

• Training dynamics differ across models. For AM and POMO, model performance and Probing Task 2 results improved fastest in the initial epochs. LEHD achieved peak Probing Task 2 performance right from the start of training.

• CVRP constraint knowledge is captured. Probing Tasks 3 and 4 confirmed that NCO models encode both the low-level linear additive relationship among node demands and the higher-level property of two nodes belonging to the same route in the optimal solution.

• Probing generalizes beyond the main setting. The authors confirm probing works on DIFUSCO, a diffusion-based NCO model; on the Asymmetric Traveling Salesman Problem (ATSP) as a non-Euclidean case study; and on the Job Shop Scheduling Problem (JSSP) with a GNN-based neural model, where precedence constraint information was captured.

• Distinct inductive biases across models. In heatmaps of node embeddings from 50 sampled instances per model, LEHD showed strong activation concentrated in fewer than 20 fixed dimensions, with absolute values often in the tens. AM and POMO showed more dispersed activation patterns, with no consistently dominant dimensions and coefficient magnitudes all below 4 in absolute value.

• Better-generalizing models reuse the same dimensions. LEHD consistently maintained the same top-2 dimensions across generalization scenarios — dimensions 31 and 69 for Probing Task 1, and 31 and 97 for Probing Task 2. AM and POMO showed varying key dimensions across generalization settings. Applying the Benjamini-Hochberg procedure to control the false discovery rate at 0.05 across the 256 tested dimensions, the key LEHD dimensions remained significant.

• Generalization failure correlates with disorganized dimensions. When AM and POMO (trained on 20-node instances) were tested on closer scales of 21, 25, and 30 nodes, both showed similar reuse of embedding dimensions on 21-node instances, supporting the idea that achievable generalization coincides with a fixed, small set of dimensions.

• Two dimensions nearly suffice for LEHD. On 100-node problems, LEHD using all 128 dimensions achieved a 0.57% optimality gap versus 0.65% with dimensions 31 and 97, 0.75% with 32 and 97, but 183.80% with 21 and 123 and 384.19% with 7 and 78. On 200, 500, and 1000 nodes, the all-128 versus 31/97 gaps were 0.86% vs 0.93%, 1.56% vs 1.81%, and 3.17% vs 3.56% respectively.

• Ablation confirms the key dimensions are causally important. Zeroing dimensions 31 and 97 collapsed LEHD's performance to a 60.43% gap versus the 0.57% original. Zeroing dimension 31 alone gave 0.90%, dimension 97 alone 1.80%, whereas zeroing non-key dimensions 98 (0.59%), 126 (0.60%), or both 98 and 126 (0.62%) had little impact.

• Key dimensions support human-readable interpretation. In a 2D plane formed by LEHD's dimensions 31 and 97 (random seed 1), the current node was node 4; a myopic distance-based decision would pick node 5, but the optimal solution picks node 3 — and in that 2D space node 4 is indeed closer to node 3. LEHD also consistently positioned the current node in the bottom-left corner, separated from other nodes, reflecting its architecture in which candidate node embeddings are re-encoded at each decision step.

• Regularization improves generalization. Training LEHD on TSP100 with a sparsity-promoting regularization term improved transfer: on TSP1000 the gap went from 3.17% (λ = 0) to 2.87% (λ = 1e-6), 2.93% (λ = 1e-5), 2.97% (λ = 1e-4), and 3.05% (λ = 1e-3); on TSP200, λ = 1e-4 gave 0.73% versus 0.86% at λ = 0; on TSPLib, λ = 1e-6 gave 4.94% versus 5.26% at λ = 0, though λ = 1e-3 degraded TSPLib performance to 8.61%.

• Probing also explains performance differences across problem sizes. LEHD achieved better probing results on larger instances (200-TSP), aligning with its stronger performance on large-scale TSP.

Methodology in Plain English

The authors borrow linear probing, a technique from computer vision and NLP: if a simple linear model trained on a frozen network's intermediate embeddings can predict some property, then that property is linearly encoded in those embeddings. Because combinatorial optimization lacks the ready-made subtasks that NLP has, the team hand-designed probing tasks with domain knowledge. For TSP they ask whether embeddings let a linear probe predict Euclidean distances between the current node and unvisited nodes (Task 1) and whether they let a probe distinguish optimal edges from greedy nearest-neighbor edges (Task 2). For CVRP they ask whether embeddings encode the additive relationship among demands (Task 3) and whether two nodes share a route in the optimal solution (Task 4).

Embeddings are extracted at points marked in the architecture diagrams — initial projections, encoder layers (Enc-l3 for AM, Enc-l6 for POMO), and decoder layers (Dec-l6 for LEHD) — for three models: AM, POMO, and LEHD, all using 128-dimensional embeddings. Probe regression and classification results are averaged over 10 evaluation runs and reported as mean ± SEM, with some probes including interaction terms.

CS-Probing extends this by inspecting the probing model's own coefficients: how large each is, and whether each is statistically significant. Because the probe input concatenates the current node's embedding (first 128 dimensions) with a candidate node's embedding (next 128 dimensions), the coefficients can be attributed to specific dimensions and roles. The authors compare coefficient patterns across models, across training runs, and across generalization settings, and apply false discovery rate control at 0.05 via the Benjamini-Hochberg procedure across the 256 tested dimensions.

They validate the findings causally with neuron ablation (zeroing dimensions) and a 2D visualization of the two most important LEHD dimensions, then test a sparsity-inducing regularization aimed at concentrating information into fewer dimensions.

Why This Matters

Impact on research. This is described by the authors as the first systematic attempt to interpret black-box NCO models. It replaces vague claims about "black-box" solvers with measurable, dimension-level evidence about what is encoded where, providing a reusable methodology (four probing tasks plus CS-Probing) and a publicly available codebase for the NCO community. The finding that two of 128 dimensions nearly suffice for LEHD, and that zeroing the wrong two collapses performance, has direct implications for compression, distillation, and pruning research and for how representation dimensionality is chosen in NCO architectures.

Real-world applications (framed as implications of interpretability; the paper does not present deployed case studies):

• Routing and logistics: explaining and auditing learned TSP or vehicle-routing solvers before they replace classical heuristics such as Concorde, LKH3, or HGS in dispatch systems.
• Engineering and scientific computing: building trust in learned optimizers deployed in domains where black-box decisions carry reliability risk.
• Model deployment and maintenance: using identified key dimensions to compress or prune solver representations for faster, lighter inference.
• Scheduling: the paper's JSSP probing results extend the interpretability approach to job shop scheduling, relevant to manufacturing and production planning.

Industry relevance. The paper explicitly notes that industries remain cautious about deploying NCO methods to replace classical operations research methods because these DL-based methods are perceived as black-box models lacking the reliability and interpretability of traditional OR approaches. The demonstrated regularization improvement — achieved with minor code modifications — shows that probing insights can feed directly back into model design and generalization, which is precisely the kind of evidence needed for adoption decisions.

Future Directions

• Investigate what additional information LEHD learns beyond myopia avoidance, since its Probing Task 2 performance peaked from the start of training; the authors suggest designing new probing tasks for this.
• Extend probing to more and different NCO model families, including deeper investigations of diffusion-based models such as DIFUSCO and other diffusion approaches.
• Explore whether embedding dimensionality within LEHD can be further improved through architectural or training adjustments, and how dimensionality should be configured across different layers and training stages.
• Pursue compressing, distilling, and pruning the representation space of NCO models, building on the result that inference with only two dimensions yields comparable results.
• Develop more probing tasks to progressively transform black-box representations into interpretable, structured forms, analogous to how DNA sequencing transformed genetics.

Target Audience

Researchers and graduate students in machine learning, interpretability, and operations research who work on neural combinatorial optimization or learned solvers; practitioners considering NCO deployment in routing, logistics, or scheduling who need evidence about reliability and generalization; and model designers interested in how representation dimensionality and sparsity affect transfer to larger problem sizes and new distributions.

Authors’ abstract

Neural combinatorial optimization (NCO) has achieved remarkable performance, yet its learned model representations and decision rationale remain a black box. This impedes both academic research and practical deployment, since researchers and stakeholders require deeper insights into NCO models. In this paper, we take the first critical step towards interpreting NCO models by investigating their representations through various probing tasks. Moreover, we introduce a novel probing tool named Coefficient Significance Probing (CS-Probing) to enable deeper analysis of NCO representations by examining the coefficients and statistical significance during probing. Extensive experiments and analysis reveal that NCO models encode low-level information essential for solution construction, while capturing high-level knowledge to facilitate better decisions. Using CS-Probing, we find that prevalent NCO models impose varying inductive biases on their learned representations, uncover direct evidence related to model generalization, and identify key embedding dimensions associated with specific knowledge. These insights can be potentially translated into practice, for example, with minor code modifications, we improve the generalization of the analyzed model. Our work represents a first systematic attempt to interpret black-box NCO models, showcasing probing as a promising tool for analyzing their internal mechanisms and revealing insights for the NCO community. The source code is publicly available.