ivanjaenm/torot-1d

ToR-OT 1D model (157.6M parameters)

Model of Learning Optimal Transport Plans Via Autoregressive Token Regression (NeurIPS 2026), used for every 1D result of the paper. An encoder reads the two histograms and the cost type; a decoder writes the optimal transport plan row by row as tokens (target-bin pointer, exponent, four mantissa digits). Code: https://github.com/vsingh-group/TOR-OT

• Weights: exponential moving average after 200000 training steps (model.safetensors, float32); architecture in config.yaml.
• Training data: ivanjaenm/torot_data_5M at revision 54892f9763ce5b9b7237a2790aed8fff85e98baf: 1D histograms uniform on the probability simplex on a regular grid of n = 5-100 bins in [0, 1], six cost types (euclidean, sqrt_euclidean, sqeuclidean, cubic_euclidean, canberra, gaussian).
• Decoding used for the paper: greedy, grammar-constrained, row by row, last row decoded by the model (last_row="model").

import numpy as np, torch, ot
from torot.evaluate import load_model
from torot.decode import predict_plans
from torot.data.cost import COST_TYPES, cost_matrix
from torot.data.supports import grid_1d

device = torch.device("cuda")
model, tok, cfg = load_model("ivanjaenm/torot-1d", device)

n = 30
x = grid_1d(n)                                           # bin centres on [0, 1], shape [n, 1]
rng = np.random.default_rng(0)
a, b = rng.dirichlet(np.ones(n)), rng.dirichlet(np.ones(n))
inst = {"n": n, "m": n, "dim": 1, "cost_type": COST_TYPES.index("sqeuclidean"),
        "coords_src": x, "coords_tgt": x, "a": a, "b": b}
plan = predict_plans(model, tok, [inst], device, last_row="model")[0]           # dense n x m plan

C = cost_matrix(x, x, "sqeuclidean")
print((plan * C).sum() / (ot.emd(a, b, C) * C).sum() - 1)                       # relative cost gap to the exact plan

Results on the corpus' test splits (weighted MAPE = the paper's MAPE)

split	instances	wMAPE (median)	wMAPE (mean)	support recall	max violation (median)
test_id	5000	0.0180	0.1746	0.903	5.3e-04
test_ood	1728	0.1144	0.4516	0.751	7.4e-04
test_extrap	1000	0.7129	0.8186	0.586	1.1e-03

baselines/: the Meta OT and UNOT runs used in the paper (one per cost, selected on validation; code in baselines/ of the repository).

Citation

@inproceedings{jaenmarquez2026torot,
  title     = {Learning Optimal Transport Plans Via Autoregressive Token Regression},
  author    = {Jaen Marquez, Ivan and Chytas, Sotirios Panagiotis and Singh, Vikas},
  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
  year      = {2026}
}