---
license: mit
library_name: torot
tags: [optimal-transport, transformer, token-regression]
datasets: [ivanjaenm/torot_data_4M_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`](https://huggingface.co/datasets/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"`).

```python
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

```bibtex
@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}
}
```