CombiLatent

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.

Full short paper: short_paper_submission_icml_2026_workshop_CombiLatent.pdf.

The Permutation Flow Shop Scheduling Problem (PFSP)

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

Phase 1 — Surrogate training

A decoder-only Transformer with relative positional embeddings is trained via supervised regression on (permutation, makespan) pairs:

$\mathcal{L}_{\text{MSE}} = \frac{1}{B}\sum_{b=1}^{B}\bigl(f_\theta(X_1^{(b)}) - y^{(b)}\bigr)^2$

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:

$\mathcal{L}_{\text{total}} = f_\theta(X_2) + \lambda \cdot \mathcal{L}_{\text{Sinkhorn}}(X_1, X_2)$

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.

Repository layout

src/
  xp0/                          # Main paper — Sections 3-4 (15-instance grid) + Appendix B (ablations)
    source/
      create_dataset.py         # PFSP instance generation + exhaustive makespan enumeration
      process_dataset.py        # Top-k filtering + train/val split + normalization
      train.py                  # GPT-style surrogate architecture + Phase 1 training loop
      recover_schedules.py      # Phase 2 + Hungarian recovery
      launch_create_dataset.py  # Grid driver — dataset generation across the 15 PFSP instances
      launch_process_dataset.py # Grid driver — top_0 … top_4 degradation across instances
      launch_train.py           # Grid driver — Phase 1 across (model_size × weight_decay)
      launch_recover_schedules.py # Grid driver — Phase 2 across all configs (incl. ablations)
      viz_train.ipynb           # Phase 1 training visualization
      viz_rs.ipynb              # Phase 2 optimization visualization
    datasets/
      exhaustive_{n}_{m}/       # Per-instance raw datasets and per-config Phase-2 outputs
      aggregated_min_rank_{1,2,3}.json          # Min-Rank-k tallies → Figure 3
      no_sinkhorn_aggregated_min_rank_*.json    # Sinkhorn ablation aggregates → Figure 5
      no_langevin_aggregated_min_rank_*.json    # Langevin ablation aggregates → Figure 6
      flowshop_transformer_heatmaps.png         # Figure 3
      no_sinkhorn_flowshop_transformer_heatmaps.png  # Figure 5
      no_langevin_flowshop_transformer_heatmaps.png  # Figure 6
    presentation/
      final_report.pdf, final_slides.pdf

  xp1/                          # Appendix D.1 — 9×6 surrogate-parameterization sweep
    1_run.py                    # Stage 1: PFSP instance generation
    2_run.py                    # Stage 2: dataset processing (top_k filtering + train/val split)
    3_run.py                    # Stage 3: Phase 1 — train surrogate for the full architectural grid
    4_run.py                    # Stage 4: Phase 2 — Sinkhorn-regularized latent optimization
    src/
      create_dataset.py, process_dataset.py, train.py, recover_schedules.py
      viz_train.ipynb, viz_rs.ipynb
    outputs/exhaustive_9_6/
      results.ipynb             # Aggregates results.json across configs → Figures 7 & 8

  xp2/                          # Appendix D.2 — data-quality study
    _run.py                     # Top-p% removal sweep + train + Phase 2 on the 9×6 instance
    outputs/exhaustive_9_6/
      results.ipynb             # → Figure 9

  xp3/                          # Exploratory: iterative surrogate refinement (not in paper)
    _run.py                     # Phase 2 captures schedules → fed back into Phase 1 training pool
    recover_schedules.py        # Phase 2 with schedule-capture instrumentation

short_paper_submission_icml_2026_workshop_CombiLatent.pdf

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:

Stage	Launcher	Module
1. Dataset generation	launch_create_dataset.py	create_dataset.py — exhaustive `(permutation, makespan)` enumeration + NEH/CDS/Palmer baselines
2. Dataset processing	launch_process_dataset.py	process_dataset.py — `top_k` filtering + train/val split + normalization
3. Phase 1 training	launch_train.py	train.py — Phase 1 over `(model_size × weight_decay)`
4. Phase 2 recovery	launch_recover_schedules.py	recover_schedules.py — Phase 2 + Hungarian recovery

Per-instance datasets and per-config Phase-2 outputs are stored under src/xp0/datasets/exhaustive_{n}_{m}/. The Min-Rank-$k$ aggregates that produced Figure 3 of the paper are in aggregated_min_rank_1.json, aggregated_min_rank_2.json, aggregated_min_rank_3.json, and the figure itself is checked in as flowshop_transformer_heatmaps.png.

Per-config training-loss curves and Phase-2 trajectories can be inspected with src/xp0/source/viz_train.ipynb and src/xp0/source/viz_rs.ipynb.

Ablation studies — Appendix B

Both ablations reuse the best config identified in Section 4 (Sm, weight decay = 10.0) and toggle a single component in the Phase 2 call:

Ablation	Toggle in `recover_schedules(...)`	Aggregated results	Paper figure
Remove Sinkhorn regularization	$\lambda = 0$	no_sinkhorn_aggregated_min_rank_{1,2,3}.json → no_sinkhorn_flowshop_transformer_heatmaps.png	Figure 5
Remove Langevin-like noise	disable periodic Gaussian injection on $X_2$	no_langevin_aggregated_min_rank_{1,2,3}.json → no_langevin_flowshop_transformer_heatmaps.png	Figure 6

Both are run through src/xp0/source/recover_schedules.py — same Phase 1 checkpoints, only the Phase 2 hyperparameters change.

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.

Driver: src/xp2/_run.py — implements its own process_dataset_pct(...) helper for the percentile-based filtering, then reuses src/xp1/src/train.py and src/xp1/src/recover_schedules.py. Aggregation and Figure 9 are produced by src/xp2/outputs/exhaustive_9_6/results.ipynb.

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:

import os
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
torch.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.

Downloads last month: -; Downloads are not tracked for this model. How to track

Inference Providers NEW

This model isn't deployed by any Inference Provider. 🙋 Ask for provider support