An Edge-Aware Graph Autoencoder Trained on Scale-Imbalanced Data for Traveling Salesman Problems

TL;DR

This work tackles the Traveling Salesman Problem (TSP) across diverse scales under scale-imbalanced data. It reframes TSP as a link-prediction task on sparse graphs and introduces EdgeGAE, combining a residual gated encoder with an edge-centered decoder to learn edge embeddings that drive solutions. To handle imbalanced data, the authors propose Random Active Sampling and demonstrate performance on a 50,000-instance benchmark spanning 50–500 cities, with tests up to 700 cities, showing competitive optimal gaps and strong generalization alongside favorable inference times. The approach offers a scalable, data-efficient neural solver for TSP and lays groundwork for extensions in semi-supervised, reinforcement, and federated settings for real-world optimization tasks.

Abstract

In recent years, there has been a notable surge in research on machine learning techniques for combinatorial optimization. It has been shown that learning-based methods outperform traditional heuristics and mathematical solvers on the Traveling Salesman Problem (TSP) in terms of both performance and computational efficiency. However, most learning-based TSP solvers are primarily designed for fixed-scale TSP instances, and also require a large number of training samples to achieve optimal performance. To fill this gap, this work proposes a data-driven graph representation learning method for solving TSPs with various numbers of cities. Specifically, we formulate the TSP as a link prediction task and propose an edge-aware graph autoencoder (EdgeGAE) model that can solve TSPs by learning from various-scale samples with an imbalanced distribution. A residual gated encoder is trained to learn latent edge embeddings, followed by an edge-centered decoder to output link predictions in an end-to-end manner. Furthermore, we introduce an active sampling strategy into the training process to improve the model's generalization capability in large-scale scenarios. To investigate the model's practical applicability, we generate a scale-imbalanced dataset comprising 50,000 TSP instances ranging from 50 to 500 cities. The experimental results demonstrate that the proposed edge-aware graph autoencoder model achieves a highly competitive performance among state-of-the-art graph learning-based approaches in solving TSPs with various scales, implying its remarkable potential in dealing with practical optimization challenges.

Figures (7)

• Figure 1: Solving TSPs as a link prediction task on graphs. In the heuristic sparsification stage, the input coordinates of $N$ cities are first converted into a complete graph where all nodes are directly connected, followed by heuristics such as $k$-nearest neighbours to get a sparse graph without redundant edges. In the link prediction stage, the sparse graph with node and edge features is input to the encoder to get latent graph embeddings via multiple layers of message passing through edges (in yellow arrows). Then the decoder performs link prediction on each edge by aggregating the embeddings of the source and target node into the edge embedding via an edge-centered message passing scheme (in green arrows). Finally, the output heatmaps are converted into feasible solutions through post-hoc graph search approaches.
• Figure 2: The residual gated encoder. (1) The input to an encoder model is a directed sparse graph, where the input node features and edge features are represented as $\alpha$ and $\beta$, respectively. In the sparse graph representation of the TSP, each node is connected to its $k$ nearest neighbours. (2) The node and edge features are mapped to high-dimensional embeddings through linear projections to enhance the expression power. (3) Multiple message-passing layers enable nodes and edges to learn information from local neighbourhoods via residual gated graph convolutions. (4) The encoder outputs the graph with latent space embeddings which can capture inherent knowledge of the graph topology for downstream tasks.
• Figure 3: The edge-centered decoder. (1) We formulate the input graph of the decoder into an edge-centered representation for a clear illustration, where each edge embedding is associated with its source and target node embeddings. (2) Three different linear projections are applied to all source nodes, target nodes, and edge embeddings, respectively. (3) Embeddings from the source and target nodes of each edge are summed up, and then aggregated into the edge embedding in an edge-centered message-passing manner. (4) The inner product result of the sigmoid-transformed node embedding and the edge embedding is passed through a multi-layer perceptron (MLP) classifier to output the probability of the edge existing in the optimal tour.
• Figure 4: Training data distributions with different sampling strategies. The random sampling retains the original distribution of the training dataset, where small-scale data is the majority and large-scale data is the minority. The active sampling strategy can change the original distribution to a balanced distribution of the number of nodes.
• Figure 5: TSP instances of different scales in the training dataset. Blue edges indicate the optimal tour obtained by the Concorde solver.