Neural Combinatorial Optimization via Latent Space Search under Sinkhorn Divergence Regularization.
CombiLatent is a general two-phase neural framework for solving NP-hard combinatorial optimization problems through continuous latent-space optimization. A Transformer surrogate is first trained to predict the objective value of any candidate solution, embedding the combinatorial structure into a differentiable latent space. A learnable solution tensor is then optimized by gradient descent against the frozen surrogate, with Sinkhorn divergence (entropic optimal transport) keeping the continuous solution anchored to a valid discrete permutation, which is finally recovered via the Hungarian algorithm.
As a first case study, we instantiate CombiLatent on the Permutation Flow Shop Scheduling Problem (PFSP) — a problem of significant industrial relevance that remains underexplored in the Neural Combinatorial Optimization (NCO) literature, which has focused overwhelmingly on routing problems (TSP, VRP). Our experiments benchmark CombiLatent against the classical NEH, CDS, and Palmer heuristics, and analyse the effects of model capacity, weight decay, and training-data quality on the learned latent space.
The Permutation Flow Shop Scheduling Problem (PFSP)
PFSP example — 3 jobs, 3 machines: two schedules of the same jobs yield makespans of 49 h vs 39 h
Given $n$ jobs and $m$ machines, each job must be processed on every machine in the same fixed order (machine 1 → machine 2 → … → machine $m$), and every machine processes jobs in the same global order — i.e. the schedule is a single permutation of the $n$ jobs shared across all machines. The processing time of job $j$ on machine $k$ is a fixed entry $p_{jk}$ of a given $n \times m$ execution-time matrix. The objective is to find the permutation that minimizes the makespan — the time at which the last job finishes on the last machine.
In the example above (3 jobs × 3 machines, times in hours), the same three jobs scheduled in two different orders produce makespans of 49 h and 39 h: a ~20% improvement obtained purely by reordering, with no change to the underlying work. The combinatorial part of the problem is exactly this choice of permutation; the search space has size $n!$ and is NP-hard.
Classical heuristics — NEH (Nawaz–Enscore–Ham, 1983), CDS (Campbell–Dudek–Smith, 1970), and Palmer (1965) — remain strong baselines in operations research and are the references CombiLatent is benchmarked against.
The Two-Phase Approach
Figure 1 — Global architecture of the CombiLatent / Flowshop Transformer pipeline
Phase 1 — Surrogate training
A decoder-only Transformer with relative positional embeddings is trained via supervised regression on (permutation, makespan) pairs:
The architectural choice (decoder-only + relative positions) is motivated by a key inductive bias: computing the PFSP makespan is essentially a dynamic-programming algorithm that processes jobs sequentially, which mirrors the autoregressive structure of decoder-only Transformers. Once trained, the model's token embedding table $X_1 \in \mathbb{R}^{n \times d}$ contains rich latent representations of the jobs. The Transformer's weights $\theta$ are then frozen.
Phase 2 — Latent schedule optimization
We introduce a new learnable tensor $X_2 \in \mathbb{R}^{n \times d}$ — one embedding per slot in the optimal schedule — initialized as a random permutation of $X_1$. We optimize $X_2$ by gradient descent against the frozen surrogate, while regularizing with Sinkhorn divergence between $X_1$ and $X_2$ to enforce that $X_2$ remains a valid permutation of real jobs:
with $\mathcal{L}{\text{Sinkhorn}}(X_1, X_2) = \min{P \in U(1_n, 1_n)} \sum_{ij} P_{ij} C_{ij} - \epsilon H(P)$, where $C_{ij} = \lVert (X_1)_i - (X_2)_j \rVert^2$ and $H(P)$ is the entropic regularizer. Langevin-like Gaussian noise is periodically injected into $X_2$ to escape local minima. Once $X_2$ converges, the Hungarian algorithm projects it back to a discrete permutation of jobs.
Note. The code under xp0/ was migrated from an earlier repository; xp1–xp3 were developed in the current repository. The four sub-trees share the same two-phase architecture but have independent code (no cross-imports between xp0/ and xp1–xp3), so each can be run in isolation.
Mapping experiments → code
Main results — Section 3 (Experimental Setup) and Section 4 (Results)
All code and results for the main paper (Sections 3-4 and the Appendix-B ablations) live under src/xp0/, migrated from an earlier repository. The numbered drivers and aggregation notebooks in xp1–xp3 cover the appendices only.
The headline grid sweeps 15 PFSP instances (jobs $n \in {7,8,9}$ × machines $m \in {2,3,4,5,6}$, execution times $\sim \mathcal{U}(0,1)$) × 5 dataset-degradation levels (top_0 … top_4, removing the 0–4 best schedules from training) × 3 model sizes (Sm ≈ 100K, Mm ≈ 800K, Bm ≈ 6M params) × 4 weight decays (0.1, 1.0, 5.0, 10.0). Performance is reported as Min Rank $k$ — the number of instances (out of 15) where CombiLatent ranks among the top $k$ against NEH, CDS, and Palmer.
The four pipeline stages each have a launcher under src/xp0/source/ that iterates the parameter grid and calls the corresponding module:
Surrogate-parameterization study — Appendix D.1 (Figures 7 & 8)
A focused sweep on the hardest single instance (9 jobs, 6 machines) over all combinations of $n_{\text{embd}} \in {6, 8, \ldots, 66}$ (divisible by $n_{\text{head}}$), $n_{\text{head}} \in {1, 2, 3, 4}$, $n_{\text{layer}} \in {1, 2, 3, 4}$ — 200 configurations in total. Each config is scored by whether Phase 2 recovers the global optimum (we use the top_0 dataset, so the optimum is present in training data).
The four numbered drivers under src/xp1/ — 1_run.py, 2_run.py, 3_run.py, 4_run.py — execute stages 1–4 of the pipeline in order, with the train/recover loops restricted to the 9×6 instance and the architectural grid above.
The aggregation notebook src/xp1/outputs/exhaustive_9_6/results.ipynb walks every train_nEmbd*_nHead*_nLayer*/recover_no_mf/results.json and produces Figure 7 (bar chart of win rate by n_embd) and Figure 8 (heatmap of win rate by (n_head, n_layer)).
Data-quality study — Appendix D.2 (Figure 9)
For each percentile $p \in {5, 10, 15, \ldots, 95}$, the top-$p$% best schedules are removed from the training data of the 9×6 instance, then the surrogate is retrained from scratch with the best architectural config from D.1 (n_embd=40, n_head=4, n_layer=2) and Phase 2 is run without Langevin dynamics to isolate the effect of data quality.
Exploratory — iterative surrogate refinement (xp3, not in paper)
src/xp3/_run.py explores a feedback loop where Phase 2's intermediate schedule captures are fed back into the Phase 1 training pool over multiple iterations. Not reported in the short paper.
Phase 2 hyperparameters (best config)
Hyperparameter
Value
Sinkhorn coefficient $\lambda$
2.0
Sinkhorn $\epsilon$
0.01
Sinkhorn iterations
40
Phase 2 gradient steps
1500 (main grid) / 2000 (Appendix D)
Langevin noise injection
every 200 steps
Hardware
single NVIDIA RTX 3070 (8GB VRAM)
Model configurations
Config
Embedding dim
Heads
Layers
Params
Sm (Small)
64
4
2
≈ 100K
Mm (Medium)
128
8
4
≈ 800K
Bm (Big)
256
16
8
≈ 6M
All models share an FFN width multiplier of 4. Phase 1 is trained for 5 epochs.
Key findings
Smaller surrogates win.Bm overfits to combinatorial artifacts; its latent manifold is too complex for Phase 2 gradient descent to navigate. Sm's limited capacity forces it to learn the broader "physics" of scheduling, giving a better inductive bias.
Weight decay is the most impactful hyperparameter. High weight decay (5.0–10.0) smooths the latent loss landscape, and its importance grows as training data degrades.
Sinkhorn is structurally necessary, not optional. Without it, $X_2$ converges to continuous vectors that do not correspond to any real job, rendering the Hungarian projection unreliable. Min Rank 1 collapses from ~2/15 to 1/15, Min Rank 3 from ~11/15 to ~6/15.
Langevin noise is complementary. Removing it drops Min Rank 3 to ~8–9/15 — meaningful but moderate, since high weight decay already provides substantial landscape smoothing.
Reproducibility note
The Appendix-D.1 grid was lost from disk and re-run after the paper was submitted. The re-run produces win-rate numbers that are close to but not bit-identical with the values printed in Figures 7 and 8 of the paper. The discrepancy is expected, not a bug, and the qualitative findings are preserved.
What happened. Both train(...) and recover_schedules(...) explicitly seed torch, numpy, and random at entry. That is necessary but not sufficient for bit-exact reproducibility on CUDA: the code does not set torch.backends.cudnn.deterministic = True, torch.use_deterministic_algorithms(True), or CUBLAS_WORKSPACE_CONFIG. Without those, several CUDA kernels (cuBLAS GEMM split-K paths, scaled-dot-product-attention backend selection, atomic-add reductions) reorder floating-point operations between runs. Each individual reordering is invisible (~1e-7), but the discrepancies accumulate through two stacked iterative optimizations:
Phase 1 runs ~3000 gradient steps. Per-step FP drift → the final surrogate weights $\theta$ sit at a slightly different point in parameter space, and the learned job embeddings $X_1$ are at slightly different coordinates.
Phase 2 then runs 2000 gradient steps against that slightly different surrogate, in a non-convex landscape with periodic Langevin noise injection. Small differences in curvature near a saddle can flip which basin of attraction Phase 2 descends into.
This exactly matches the pattern observed in the re-run: robust cells reproduce, borderline cells drift by 10–20 percentage points.
Comparing the re-run to the paper. Side-by-side win rates (%) for the (n_head, n_layer) heatmap (Figure 8):
(n_head, n_layer)
(1,1)
(1,2)
(1,3)
(1,4)
(2,1)
(2,2)
(2,3)
(2,4)
(3,1)
(3,2)
(3,3)
(3,4)
(4,1)
(4,2)
(4,3)
(4,4)
Paper
29
43
57
71
43
14
29
57
24
57
48
71
33
73
47
67
Re-run
29
43
71
71
43
14
43
71
33
62
48
71
40
80
47
80
Δ
0
0
+14
0
0
0
+14
+14
+9
+5
0
0
+7
+7
0
+13
All three qualitative conclusions the paper draws from this study survive:
n_embd=40 is a sweet spot. Paper: 100% win rate across all (n_head, n_layer). Re-run: 100% — identical.
The (n_head=2, n_layer=2) valley. Paper: 14% (the lowest cell). Re-run: 14% — identical.
The "asymmetric structure" finding. The paper observes two close-best cells: high $n_{\text{head}}$ × low $n_{\text{layer}}$ (paper: (4,2) = 73%) and high $n_{\text{layer}}$ × low $n_{\text{head}}$ (paper: (1,4) and (3,4) = 71%). Both peaks are still present in the re-run — only their relative heights shift by a few percentage points.
Overall win rate. Paper-implied total ≈ 51% (sum of cells / 200). Re-run: 54.5%. Same order of magnitude.
The re-run is also subject to environment drift since the original training: PyTorch / CUDA / driver versions may have changed kernel selection (e.g., SDPA's flash-attention vs math fallback) between the original run and the recovery, on top of the per-run kernel non-determinism above.
For bit-exact future runs, add at the top of train(...) and recover_schedules(...), right before the seed calls:
python
1import os
2os.environ["CUBLAS_WORKSPACE_CONFIG"]=":4096:8"3torch.backends.cudnn.deterministic =True4torch.backends.cudnn.benchmark =False5torch.use_deterministic_algorithms(True, warn_only=True)
Even with those, cross-machine reproducibility (different GPU model, different CUDA version) is not guaranteed by PyTorch.
Limitations / future work
Instance-specific retraining — the surrogate must currently be retrained per PFSP instance. Conditioning it on an MLP-encoded representation of the execution-time matrix would enable cross-instance generalization.
Latent landscape non-convexity — Phase 2 gradient descent remains prone to local minima. Principled non-convex optimizers (e.g. a latent "Tabu search" that penalizes proximity to detected local minima) are a promising direction.
Experimental scale — restricted to $n \in [7,9]$, $m \in [2,6]$ due to the cost of exhaustive enumeration. Scaling to Taillard-benchmark sizes ($n \in [20, 500]$) is the most important direction for future work.