Generalize a Small Pre-trained Model to Arbitrarily Large TSP Instances

In plain words

The travelling salesman problem asks for the shortest round trip through a set of cities. Exact solvers and strong classical heuristics exist, but they are intricate and built from hand-written rules. Learned solvers learn from examples, yet a model trained on one problem size usually does poorly on other sizes, and collecting optimal answers for big problems is costly.

The idea here is to train one small model on small problems and then reuse it on huge maps. For a large map, the method repeatedly takes a small neighbourhood of cities, asks the model which connections look promising inside it, and averages these local answers into one map of likely connections. A search then improves a tour, favouring the connections that the map rates highly.

The model never has to be retrained for larger sizes. The method targets flat, two-dimensional maps; extending it to other kinds of distance is left for future work.

TL;DR

• One fixed-size graph-attention network (Att-GCRN) is trained with Concorde-optimal tours as labels §3.1.
• For an instance of any size, it samples overlapping sub-graphs of the trained size, predicts edge probabilities on each, and averages them into one heat map §3.3 fig. 2.
• Monte Carlo tree search over k-opt moves, guided by the heat map, refines the tour §3.4 fig. 3.
• Author-reported gaps: 0.0000% / 0.0145% / 0.0370% on 20 / 50 / 100 cities tab. 1; 0.8844% / 2.5365% / 3.2238% on 200 / 500 / 1000 cities tab. 2; 4.3902% vs LKH3 on 10,000 cities, 16 instances tab. 3.
• Code (MIT, linked from the paper) includes a Python heat-map pipeline and C++/CUDA search. Several constants and formulas differ from the paper; see Paper vs code.

Problem

• The travelling salesman problem (TSP): find the cheapest tour that visits every city once and returns to the start. It is NP-hard §1.
• Exact solvers (Concorde) and strong heuristics (LKH3) work well but are complicated and rely on expert rules, which makes them hard to transfer to other combinatorial problems §1.
• Supervised learned models fit a fixed size. Their quality can drop sharply on instances of a different size, and training needs many pre-computed optimal or near-optimal tours, which is unaffordable for large instances §1 §3.3.
• Scope: two-dimensional Euclidean TSP, with coordinates in the unit square §3.1. Non-Euclidean instances are future work §5.

Method

Three stages §3.2 fig. 1:

1. Offline learning of the Att-GCRN model on instances of a fixed size m, labelled with Concorde tours §3.1.
2. Heat map for the large instance: sample overlapping sub-graphs of size m, rescale each one, predict its heat map with the model, then merge the predictions per edge §3.3 fig. 2.
3. Online search: start from a constructed tour, improve it with 2-opt moves restricted to promising edges, then run MCTS over larger k-opt actions with a weight matrix W initialised from the heat map §3.4 fig. 3.

Heat-map model §3.1

• Att-GCRN is a graph convolutional residual network with an attention mechanism. The paper says its architecture is described in the full version, not in this paper §3.1.
• Input: the m city coordinates. Output: for each edge (i, j), a probability P_ij in [0, 1] that the edge belongs to the optimal tour. The heat map is an n × n matrix §3.1.
• Training set: 990,000 random instances of m cities, labelled with Concorde solutions §3.1.

Sub-graph sampling and conversion §3.3.1 §3.3.2

• Visit counts: O_i for each city and O_ij for each edge, all starting at 0 §3.3.1.
• Repeat: pick the centre as a city with the smallest O_i, ties broken at random; take the centre and its m − 1 nearest cities as the sub-graph; increase O_i and O_ij for its members §3.3.1.
• Stop when the smallest O_i reaches the lower bound ω §3.3.1. Sub-graphs may overlap §3.3.1.
• Conversion: with x_min, x_max, y_min, y_max over the sub-graph, s = 1 / max(x_max − x_min, y_max − y_min), and each coordinate becomes s × (x − x_min) and s × (y − y_min) eq. 1 §3.3.2.

Merging sub heat maps §3.3.3 §3.3.4

• Run the model on each converted sub-graph to get its heat map §3.3.3.
• Merge: P_ij = (1 / O_ij) × Σ over the sub-graphs l that contain edge (i, j) of P''_ij(l) eq. 2 §3.3.4.
• Edges with P_ij < 10^-4 are unpromising and are removed from the search §3.3.4.

# Heat map for an n-city instance  [§3.2] [fig. 2]
# Input:  coords (n, 2) in the unit square; model trained on m-city tours
# Output: heat map P (n, n), each entry in [0, 1]

def build_heat_map(coords, model, m, omega):
    # Visit counts: O_i per city and O_ij per edge  [§3.3.1]
    city_counts = zeros(n)
    edge_counts = zeros(n, n)
    edge_sums = zeros(n, n)
    # Keep sampling until the least-visited city has omega visits  [§3.3.1]
    while min(city_counts) < omega:
        # Centre: a least-visited city, ties broken at random  [§3.3.1]
        center = random_choice(where(city_counts == min(city_counts)))
        # Sub-graph: the centre and its m - 1 nearest cities  [§3.3.1]
        subset = [center] + nearest_cities(coords, center, m - 1)
        city_counts[subset] += 1
        edge_counts[subset, subset] += 1
        # Rescale the sub-graph to the unit square  [§3.3.2] [eq. 1]
        local = rescale(coords[subset])
        # Model gives a probability for every edge of the sub-graph  [§3.3.3]
        local_probs = model(local)
        # Add the probabilities onto the original edge indices  [§3.3.4]
        edge_sums[subset, subset] += local_probs
    # Average each edge over the sub-graphs that contained it  [§3.3.4] [eq. 2]
    return edge_sums / edge_counts

Search on the heat map §3.4

States and actions §3.4.1

• A state is a full tour, a permutation of all n cities.
• An action is a k-opt move written as a_1, b_1, ..., a_k, b_k, a_{k+1} with a_{k+1} = a_1. It deletes edges (a_i, b_i) and adds edges (b_i, a_{i+1}). Given a_i, b_i is determined, so an action needs only k choices §3.4.1.
• Change in length: Δ = Σ d(b_i, a_{i+1}) − Σ d(a_i, b_i) eq. 3 §3.4.1. The action improves the tour when Δ < 0 §3.4.1.
• Actions that add an edge with P < 10^-4 are discarded §3.4.1.

State initialisation §3.4.2

• Start at an arbitrary city and repeatedly move to an unvisited city, chosen with probability proportional to exp(P_{π_i j}) §3.4.2.

Small neighbourhood, k = 2 §3.4.3

• Examine promising actions with k = 2 one by one, apply the first one that improves the tour, and repeat until none does §3.4.3.

Enlarged neighbourhood, MCTS §3.4.4 fig. 3

• Initialisation, once for the whole run: W_ij = 100 × P_ij, Q_ij = 0 (times edge (i, j) was chosen in simulations), M = 0 (actions simulated) §3.4.4.
• Simulation: choose a_1 at random, then b_1 is fixed by a_1. For each later step, the potential of an edge from b_i to j is Z = W_{b_i j} / Ω_{b_i} + α × sqrt(ln(M + 1) / (Q_{b_i j} + 1)), where Ω_{b_i} is the mean of W over all edges from b_i eq. 4 §3.4.4.
• Choosing a_{i+1}: if closing the loop (linking b_i to a_1) gives an improving action, or i ≥ 10, set a_{i+1} = a_1. Otherwise the candidates X are the cities j with W_{b_i j} ≥ 1, excluding a_1 and the city already linked to b_i. Pick j from X with probability Z_{b_i j} / Σ_{l in X} Z_{b_i l} eq. 5 §3.4.4.
• Sampling continues until an improving action is found or the pool reaches H actions (H = 10 n) §3.4.4 §4.2.
• Selection: if an improving action is found, apply it. If the pool has none, the search jumps to a new random state built as in state initialisation §3.4.4.
• Back-propagation: each examined action increases M by 1, and each edge (b_i, a_{i+1}) of it increases Q by 1. When a state improves, every edge (b_i, a_{i+1}) of the applied action increases by β × (exp((L(π) − L(π*)) / L(π)) − 1) eq. 6 §3.4.4. W and Q stay symmetric §3.4.4.

Termination §3.4.5

• Simulation, selection and back-propagation repeat until no improving action is in the pool. Then the search jumps to a new state and restarts the small and enlarged neighbourhood searches. The process ends when the time budget T is used up, and the best tour found is returned §3.4.5.

# Potential of edge (b, j) when sampling an action  [§3.4.4] [eq. 4]
# Input: W, Q (n, n) matrices; M scalar; alpha = 1  [§4.2]

def potential(W, Q, M, b, j):
    # Average weight of the edges leaving b  [§3.4.4]
    avg = mean(W[b, others(b)])
    # Exploitation term plus exploration term  [eq. 4]
    return W[b, j] / avg + alpha * sqrt(log(M + 1) / (Q[b, j] + 1))

# Choose a_{i+1} to link to b_i  [§3.4.4] [eq. 5]
# Input: b (current city), a_1 (start city), depth i, W, Q, M

def choose_next(b, a_1, depth, W, Q, M):
    # Close the loop at depth 10, or when closing improves  [§3.4.4]
    if depth >= 10 or closing_is_improving(b, a_1):
        return a_1
    # Candidates: W >= 1, except a_1 and the city linked to b  [§3.4.4]
    candidates = [
        j for j in cities
        if W[b, j] >= 1 and j not in (a_1, linked(b))
    ]
    # Probability proportional to the potential  [eq. 5]
    scores = [potential(W, Q, M, b, j) for j in candidates]
    return random_choice(candidates, weights=scores)

# Reward the edges of an improving action  [§3.4.4] [eq. 6]
# Input: pairs = [(b_i, a_{i+1}), ...] of the action; beta = 10  [§4.2]

def reward(W, pairs, length_before, length_after):
    # Relative gain, positive when the tour got shorter  [eq. 6]
    gain = (length_before - length_after) / length_before
    for b, a_next in pairs:
        # Same increase for every edge of the action  [§3.4.4] [eq. 6]
        W[b, a_next] += beta * (exp(gain) - 1)
        # W is symmetric  [§3.4.4]
        W[a_next, b] = W[b, a_next]

# One MCTS round: sample until an action improves, then apply it  [§3.4.4]
# W, Q and M are shared state, updated in place across rounds  [§3.4.4]

def mcts_round(tour, W, Q, M):
    while True:
        # Sampling pool of at most H = 10 n actions  [§3.4.4] [§4.2]
        action = None
        for _ in range(H):
            # Build one action from a random start; updates Q and M  [§3.4.4]
            candidate = sample_action(tour, W, Q, M)
            M += 1
            if improves(tour, candidate):
                action = candidate
                break
        # No improving action in the pool: the round ends  [§3.4.4]
        if action is None:
            return tour
        before = length(tour)
        tour = apply(tour, action)
        reward(W, pairs_of(action), before, length(tour))

# Outer loop over states until the time budget T is spent  [§3.4] [§3.4.5]
# Input: heat map P, budget T in seconds
# Output: best tour found

def solve(P, T):
    # Initial tour from the state-initialisation rule  [§3.4.2]
    tour = construct_tour(P)
    # Created once: W = 100 * P, Q = 0, M = 0  [§3.4.4]
    W, Q, M = 100 * P, zeros(n, n), 0
    best = tour
    while elapsed() < T:
        # Small neighbourhood: k = 2 moves until none improves  [§3.4.3]
        tour = two_opt_until_no_gain(tour, P)
        # Enlarged neighbourhood: MCTS rounds until none improves  [§3.4.4]
        tour = mcts_round(tour, W, Q, M)
        best = shorter(best, tour)
        # Jump to a new state and repeat  [§3.4.4] [§3.4.5]
        tour = construct_tour(P)
    return best

Hyperparameters

Results

Author-reported. Setup: test sets of random instances, Concorde as the reference for sets 1 and 2 and LKH3 for the 10,000-city set; this method runs on one GTX 1080 Ti with the search in C++; baselines re-run from their public code and pre-trained models §4 §4.1 §4.3.

Set 1: 10,000 instances each for 20, 50 and 100 cities tab. 1

Baseline times in blue were adapted to lengthen their runtime, so timings are not strictly comparable §4.3 tab. 1.

Set 2: 128 instances each for 200, 500 and 1000 cities tab. 2

Reference: LKH3 gaps are 0.0040%, 0.0029% and 0.0036%. Gurobi fails to finish at 1000 cities tab. 2 §4.4.

10,000 cities: 16 instances, gap against LKH3 tab. 3

Three learned baselines of Joshi et al. ran out of memory, and Concorde, Gurobi and two GAT models of Deudon et al. exceeded the five-hour limit on these instances §4.4.

Ablation: heat map removed (every edge gets equal probability, search only) tab. 4

The paper concludes that the heat map is what keeps the search near optimal §4.5.

Paper vs code

Repository: https://github.com/Spider-scnu/TSP at 729c52e, MIT licence LICENSE. The paper links it in a footnote §4. It has Python notebooks for the model and heat maps, a C++ MCTS with a CPU version and a CUDA version. The root README presents the CUDA build as the main path README.md. The paper says the MCTS is in C++ §4.

Heat-map model. The paper gives no architecture or training values beyond the sizes §3.1 §4.2. The repository's configs for the 20-city and 50-city models: hidden size 300, 30 GCN layers, 3 MLP layers, mean aggregation, batch size 20, 500 batches per epoch, up to 1500 epochs, learning rate 0.001 with decay 1.01, focal loss (gamma 1, 10 negatives), augmentation probability 1.0, beam size 1280 tsp20.json tsp50.json. The 50-city config uses 20 neighbours where the 20-city one uses 19 tsp50.json. Trained checkpoints are not in the repository; the README links them externally README.md. Whether the released weights match these configs was not checked.

Input graph. Each instance becomes a k-nearest-neighbour graph with num_neighbors neighbours google_tsp_reader.py.

Sub-graph sampling (differs). The paper stops when the least-covered city reaches ω = 5 §3.3.1. The notebooks instead stop when the number of sub-graphs reaches ceil(n / (K + 1) × 5) and every city has been covered at least once test-200-500-1000-10000.ipynb [#L167] [#L208-L209]. Nothing enforces a minimum of five visits per city. For 20 cities the threshold is 1 test-20-50-100.ipynb. After the first covering pass, sub-graphs are drawn at random from a larger neighbourhood of size K_expand rather than the K nearest test-200-500-1000-10000.ipynb. Values: K = 19 with K_expand = 29 or 39 in the 20-city test notebook test-20-50-100.ipynb [#L408] [#L517]; K = 49 with K_expand = 79 (200 cities), 99 (500 and 1000 cities) and 149 (10,000 cities) in the other notebook test-200-500-1000-10000.ipynb [#L404] [#L513] [#L622]. None of these K_expand values appear in the paper.

Centre choice and conversion (matches). The centre is chosen at random among the least-covered cities test-200-500-1000-10000.ipynb, and the first centre is city 0 [#L279]. Rescaling matches eq. 1 test-200-500-1000-10000.ipynb.

Merging (matches). Each edge is divided by its visit count, with 1e-8 added to avoid division by zero test-200-500-1000-10000.ipynb (same line in the 20-city notebook).

Potential, eq. 4 (CPU matches, GPU differs). The CPU code computes W/Ω + α·sqrt(ln(M+1) / (ln 2.718 × (Q+1))) TSP_MCTS.h. ln 2.718 is about 1, so this is eq. 4. The CUDA kernel computes sqrt((M+1) / (0.434 × (Q+1))) with no logarithm of M kernel.cu. The CUDA version is a different exploration term. The paper does not say which version produced its numbers.

Candidate set X (differs). The paper uses W_{b_i j} ≥ 1 §3.4.4. The CPU code keeps city j only if the eq. 4 potential is at least 1 TSP_MCTS.h. It also skips the start city and the current successor [#L52-L55].

Sampling mechanics (not in paper). Probabilities are converted to integer buckets out of 1000 by truncation, and the last bucket is set to 1000 TSP_MCTS.h [#L84]. A candidate whose share is below 0.1% gets an empty bucket and is never chosen, except that the last candidate absorbs the rounding slack and can be chosen too often.

Depth limit. The code stops extending an action when the pair count exceeds Max_Depth = 10, or when the improvement is positive TSP_MCTS.h, and also when no candidate remains [#L127-L129]. The paper's "i ≥ 10" was not checked for an off-by-one.

Best-prefix truncation (not in paper). After sampling, the code keeps only the prefix of the action with the highest real gain TSP_MCTS.h.

Back-propagation, eq. 6 (matches). The code applies β·(2.718^(Δ/L) − 1) to every edge of an improving action, with 2.718 in place of e TSP_MCTS.h. Visit counts Q and M are updated per examined action [#L132-L133]; 2-opt moves also update them (below).

2-opt updates the MCTS state (not in paper). Each 2-opt move evaluated in the local search increases the visit counts and M TSP_2Opt.h. Each applied 2-opt move also increases W on four edges by β·(2.718^(Δ/L) − 1) TSP_2Opt.h. The paper's local search (§3.4.3) describes no such update.

Initial tour (differs). The paper weights candidates by exp(P_{π_i j}) §3.4.2. The CPU code weights them by exp(W), and W = 100 × P at that point because MCTS_Init runs first TSP_Init.h [#L129] TSP_Markov_Decision.h. The CUDA version picks uniformly at random among unvisited cities TSP_Init.h.

Time budget T (differs). The paper gives T = 10 n ms for set 1 and 40 n ms for set 2 §4.2. The loop runs while elapsed time is below Param_T × n seconds TSP_MCTS.h. The values are: CPU 0.10 for 20 to 1000 cities and 10,000 cities TSP_IO.h (100 n ms); CUDA 0.01 for 20 to 100 cities TSP_IO.h (10 n ms, matching set 1); CUDA 0.006 for 200 to 1000 cities and 10,000 cities TSP_IO.h TSP_IO.h (6 n ms). No 40 n ms setting was found.

Constants that match. H = 10 TSP_IO.h, α = 1 [#L23], β = 10 [#L24], maximum depth 10 [#L20], and W initialised as 100 × P TSP_MCTS.h.

Unpromising-edge cut. The 10^-4 threshold appears in the candidate builder TSP_Basic_Functions.h, not in the merge step.

Sub-graph sizes (matches). The 20-city model is used for set 1 and the 50-city model for set 2 test-20-50-100.ipynb test-200-500-1000-10000.ipynb.

Baselines. Baseline numbers come from re-running public code, as the paper says §4. These were not checked.

Pitfalls

• Unstated model details. The architecture, optimiser, learning rate, loss and epochs are not in this paper, so a faithful reimplementation depends on the repository's configs. Those configs use focal loss and augmentation tsp20.json. The repository also ships a "full version" PDF. This card does not use it, and its licence was not checked.
• Sub-graph sampling is the main unstated design. The paper's stop rule (ω = 5 per city) is not what the notebooks do. K and K_expand values differ by instance size. Results depend on these choices.
• Division by the average weight. Eq. 4 divides by Ω_{b_i}. If all W values from a city are 0, the potential is undefined. The paper does not guard this, and the code divides without a check TSP_MCTS.h [#L40].
• No clipping in the weight update. Eq. 6 is exponential in the relative gain, and the code does not clip it TSP_MCTS.h. W can therefore grow large, which matters because the candidate test is W ≥ 1 in the paper.
• Candidate test. The paper's W ≥ 1 and the code's potential ≥ 1 select different sets. Pick one and state it in your implementation.
• Truncated probabilities. Integer buckets out of 1000 silently drop small candidates, and the last candidate absorbs the rounding error (see Paper vs code).
• Memory grows with n². The CPU code allocates double-precision n × n matrices for the heat map, and the search also keeps W and Q at that size TSP_IO.h. At 10,000 cities, one such matrix of 8-byte values is 10^8 × 8 bytes, about 0.8 GB. The paper reports memory failures for baselines at this size but not for its own method §4.4.
• Wall-clock budget. The time budget T is wall-clock, so results depend on hardware. The time column mixes heat-map time and search time tab. 1.
• Timing comparisons. Some baseline times were adapted to lengthen runtime, and non-learned methods ran on a different CPU §4 §4.3 tab. 1. Times across rows are not strictly comparable.
• Ablation table mixes references. Table 4's header says "Opt. Gap", but its 10,000-city row is measured against LKH3, not an optimum tab. 4 tab. 3.
• Inconsistent timings. Table 2 gives 20.62 s and 31.17 s for 200 and 500 cities. Table 4 gives 20.63 s and 31.73 s for the same runs. The paper does not explain this tab. 2 tab. 4.
• Small samples at scale. The 10,000-city result rests on 16 instances, and the reference is LKH3, a heuristic, not an optimum tab. 3 §4.4.
• Section numbers. The compiled PDF hides section numbers (\setcounter{secnumdepth}{0}). § references here count sections in source order: 1 Introduction, 2 Related works, 3 Methods, 4 Experiments, 5 Conclusions [raw LaTeX source].

Lineage

• Builds on: Joshi et al. 2019, a graph convolutional network for TSP (arXiv:1906.01227). The paper says its Att-GCRN is "somewhat similar" to that network and presents the sampling, conversion and merging steps as new §1.
• Compared against: Concorde, Gurobi and LKH3 §4.1; GAT models of Kool et al. (ICLR 2019) and Deudon et al. (2018); GCN of Joshi et al. §4.3 §4.4. Two MCTS-based TSP papers, Shimomura and Takashima (2016) and Xing and Tu (2020), were compared only through their reported numbers §4.4.
• Superseded by: not known from content in this wiki.