ivanjaenm/torot-2d

ToR-OT 2D model (157.6M parameters)

Model of Learning Optimal Transport Plans Via Autoregressive Token Regression (NeurIPS 2026), used for the 2D results of the paper (Fig. 4C-D, Fig. 5, App. F). 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 620000 training steps (model.safetensors, float32); architecture in config.yaml.
• Training data: ivanjaenm/torot_data_4M_2d at revision 1ff4a4a6329fa5c822f77d92b1c2ae850afedd9b: 2D histograms on k x k grids in [0, 1]^2 (k = 3-10, n = 9-100 bins), six cost types; half of the marginals have empty bins (10-70 % of the bins set to zero).
• 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_2d

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

k = 6
x = grid_2d(k)                                           # k*k bin centres on [0, 1]^2, shape [k*k, 2]
n = k * k
rng = np.random.default_rng(0)
a, b = rng.dirichlet(np.ones(n)), rng.dirichlet(np.ones(n))
inst = {"n": n, "m": n, "dim": 2, "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

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}
}