ToR-OT 2D model, best checkpoint so far (step 560000)

README.md

@@ -0,0 +1,55 @@
+---
+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 560000 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, Takis and Singh, Vikas},
+  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
+  year      = {2026}
+}
+```

SHA256SUMS

@@ -0,0 +1,45 @@
+fd3ecdce52c9f1013cd96ef66351bf1983fe30b54b9826fe21e3c02b1fc32ebe  README.md
+d62b8fb9d17855d6fca82c07de68ad21b5c932be9e03802f7f79d37cf638058e  baselines/metaot/canberra_eps0.01/config.json
+3c3d6856bd68e68b2b4ab1c5a798353e14b70fc01c089198d53f69714bc54f79  baselines/metaot/canberra_eps0.01/metrics.json
+e080d87e7b96ec52a476d78a95cbe5d5c8df533f3a819d4a9b4d443345b23d63  baselines/metaot/canberra_eps0.01/params.msgpack
+4525b3950937c1402345b8d7e12284792eb74e0eb31c06d7ceb1d18671ade71e  baselines/metaot/cubic_euclidean_eps0.001/config.json
+95ba524d3b3bb2c61a3abca70ed85bb8c2b68c84fb4520e2ff363c5d6c0aa047  baselines/metaot/cubic_euclidean_eps0.001/metrics.json
+dabb07a7b2f614b3c8dc14fcdd5a51b44d8820ac04fdee0087908be5b06460d9  baselines/metaot/cubic_euclidean_eps0.001/params.msgpack
+a7b7f1a2ae6a457a7514f3721811fbccf17ed4052a21a7dd867baf2162802a21  baselines/metaot/euclidean_eps0.01/config.json
+ba0097bea24749938e56918ea78b88a9d8da64dc4679d00913b916ed9bfb5012  baselines/metaot/euclidean_eps0.01/metrics.json
+90ed21e379ece73534f5eb030bd17fcb4f80711aa6ce04abde490cc3617e67bc  baselines/metaot/euclidean_eps0.01/params.msgpack
+aa2429c725561c0cf1c3d983969160975bb2e76f4d678cb46081b50fb3f7d586  baselines/metaot/gaussian_eps0.01/config.json
+26f4bc9d39c2def42835ae6a809f0c65e27efa3f05bc6a355d076669f0a658fb  baselines/metaot/gaussian_eps0.01/metrics.json
+03998a558012f61ca2513f809d9f4f116fe33370440ab4600509a33362e81b0a  baselines/metaot/gaussian_eps0.01/params.msgpack
+6cc68cb2753df240e8e95480f65587841c28387c703365fbf4c1a8c0d51a0c51  baselines/metaot/sqeuclidean_eps0.01/config.json
+904f0b16c08b6bf45ae2994df9bc11432931782f4453a07fdcb99b67cc1cb38f  baselines/metaot/sqeuclidean_eps0.01/metrics.json
+9cbf8a93696f96dcbd934d257d97f3ccbd3a22f40ec8b4dfe11fb1e4573e3872  baselines/metaot/sqeuclidean_eps0.01/params.msgpack
+b6ceebc1fa375e61b2947d26bf026bedc9e680b0780e40ef644f41b1a97f8a23  baselines/metaot/sqrt_euclidean_eps0.01/config.json
+755cf55b289009b1cc0a72755c7a773fe36ded274888a82218f6584b8619e3b7  baselines/metaot/sqrt_euclidean_eps0.01/metrics.json
+12de5d3b6b2c5a21bf3f82d644697011b99d71f4254d26e64ca2872ac8f98205  baselines/metaot/sqrt_euclidean_eps0.01/params.msgpack
+9084a19670ef2a9619d315a165c6afa427f5cd982693e269f14d8be3664c512e  baselines/unot/canberra_eps0.01/config.json
+e3d60a959a1da9310f9d4c5f4df82b55098fdf8be48d82babad12be06375d032  baselines/unot/canberra_eps0.01/generator.pt
+5d26e1aa6125761477334b9eb4a1c06f0fc1c1a00ee388f2dfec56c40781b30f  baselines/unot/canberra_eps0.01/metrics.json
+0b4fcbea2e5ba01bc960a5fea9c9dacee9588da2228de3dc8e5b705b790cf07e  baselines/unot/canberra_eps0.01/predictor.pt
+84102fd80db8f5ca5f480fbbb0f34fca55e6ac9f490a327863f01957601393fe  baselines/unot/cubic_euclidean_eps0.001/config.json
+2a189fe332ab67c3c9ef4fe177f6d21491c0d6bb39c4906513e32af4a8f66bea  baselines/unot/cubic_euclidean_eps0.001/generator.pt
+d478add85299c1206b9dcc5c077b76e52a408d96fd9781427155a2c1bc2c5659  baselines/unot/cubic_euclidean_eps0.001/metrics.json
+3c974ffc49bf17f273f474aadd6a7634214261b2896d5516ce33ed32c65dc1bd  baselines/unot/cubic_euclidean_eps0.001/predictor.pt
+4849b229df30f702d8a5aa0d42d6eaf995c979ce0046cea7bedc1db274b48b70  baselines/unot/euclidean_eps0.01/config.json
+beb27a049f735baf49f85aebe9e4734941595dbbc301ac2db3e89cd46a396130  baselines/unot/euclidean_eps0.01/generator.pt
+2bdf880e85df7557798d459ca1fb4b5cc98bc84f21c602a1df14c34a64782739  baselines/unot/euclidean_eps0.01/metrics.json
+e38287705ccb366e2682f264ef1023633ece05e5e319ffe6b0df90084b102223  baselines/unot/euclidean_eps0.01/predictor.pt
+c74d75b88102a3ad8d5d17c6d8ee5a453d31f29c3355d0f69a1d7e0cbb838e76  baselines/unot/gaussian_eps0.01_w144/config.json
+d1ac19634c248e2daef3c212ff7c4984868a2d2e47fa9500732d45326a3df219  baselines/unot/gaussian_eps0.01_w144/generator.pt
+c5c8e366e19891490e1223a3aa7a1846b48d8faa8e0a6147a02b462cbf52b579  baselines/unot/gaussian_eps0.01_w144/metrics.json
+daeb29419504918302b5623684e9dd3645239e6dfcfa2155f3bdb721ca18fa4b  baselines/unot/gaussian_eps0.01_w144/predictor.pt
+00f4f443220390434d2a994bb629c194450fd42366d355bd60e6d935dcdae134  baselines/unot/sqeuclidean_eps0.01_w144/config.json
+9fc5bbd5c191063a7f5ee247c2ca499b58cebcf466d5bba2cf6af06d5161624a  baselines/unot/sqeuclidean_eps0.01_w144/generator.pt
+215052692ccf920fdded52898ce000ae701389b376fd60c7b42a942c00491f9e  baselines/unot/sqeuclidean_eps0.01_w144/metrics.json
+9788fcdfff26b79823071cdb3e75d5c8bcab98cca86d9d575152d198db3bcedb  baselines/unot/sqeuclidean_eps0.01_w144/predictor.pt
+b23627da5507bfc60879af51d8d9ca6979835028a317cc6d2032d3cfab04adca  baselines/unot/sqrt_euclidean_eps0.01/config.json
+8c031c18f0ec18c52b070eb324a85d845148edaad7f21eb0bb91385ce19f8fe1  baselines/unot/sqrt_euclidean_eps0.01/generator.pt
+c98d318521b7e34ed204be16dc03eb0ee86c5346df22857a0824e1ee5c8eb65d  baselines/unot/sqrt_euclidean_eps0.01/metrics.json
+b2ec80b1427c1f07f99ed361688067b1304c5241cdd7ce470038087007a27b75  baselines/unot/sqrt_euclidean_eps0.01/predictor.pt
+c779e34fe20e08a0ca902ab484feab1d3ce5278896e39ac72462fd10d145f15d  config.yaml
+99bfdb2eedf149f2be416a4f69ad4ce37ed6b64a252c474b2cf95191fd6c92f5  model.safetensors

baselines/metaot/canberra_eps0.01/config.json

@@ -0,0 +1,47 @@
+{
+  "method": "metaot",
+  "cost": "canberra",
+  "eps": 0.01,
+  "n": 36,
+  "dim": 2,
+  "seed": 0,
+  "train_marginals": "uniform on the simplex (Dirichlet(1)), as ToR-OT's training data",
+  "official_config": {
+    "seed": 0,
+    "data": "random",
+    "potential_model": {
+      "_target_": "meta_ot.models.PotentialMLP",
+      "n_hidden": 1024,
+      "n_hidden_layer": 3
+    },
+    "batch_size": 128,
+    "num_train_iter": 50000,
+    "potential_model_mlp": {
+      "_target_": "meta_ot.models.PotentialMLP",
+      "n_hidden": 1024,
+      "n_hidden_layer": 3
+    },
+    "optim": {
+      "_target_": "optax.adam",
+      "learning_rate": 0.001
+    },
+    "max_grad_norm": 1.0
+  },
+  "vendored_commit": "7dcc9a4c25b0a57626d3fa26028c6633ad19da50",
+  "versions": {
+    "jax": "0.4.30",
+    "jaxlib": "0.4.30",
+    "flax": "0.8.5",
+    "optax": "0.2.2",
+    "numpy": "1.26.4",
+    "ott-jax": "0.2.9"
+  },
+  "device": "cuda:0",
+  "provenance": {
+    "tor_ot_commit": "b9c045807b7fbee37ab22a523e75bc3bbc076169",
+    "tor_ot_dirty": true,
+    "python": "3.10.22",
+    "executable": "/home/ivan/anaconda3/envs/env_metaot/bin/python",
+    "platform": "Linux-6.8.0-138-generic-x86_64-with-glibc2.35"
+  }
+}

baselines/metaot/canberra_eps0.01/metrics.json

@@ -0,0 +1,15 @@
+{
+  "validation": {
+    "set": "val_simplex2d_canberra",
+    "instances": 500,
+    "paper_mape_mean": 0.48757729823236323,
+    "paper_mape_median": 0.48294072325959986,
+    "row_violation_l1_mean": 0.25674583064459744,
+    "col_violation_l1_mean": 1.194925654429723e-15,
+    "abs_cost_gap_mean": 0.010577336395573882,
+    "finite_fraction": 1.0
+  },
+  "train_seconds": 4130.631784915924,
+  "parameters": 2210852,
+  "iterations": 50000
+}

baselines/metaot/canberra_eps0.01/params.msgpack

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e080d87e7b96ec52a476d78a95cbe5d5c8df533f3a819d4a9b4d443345b23d63
+size 8843673

baselines/metaot/cubic_euclidean_eps0.001/config.json

@@ -0,0 +1,47 @@
+{
+  "method": "metaot",
+  "cost": "cubic_euclidean",
+  "eps": 0.001,
+  "n": 36,
+  "dim": 2,
+  "seed": 0,
+  "train_marginals": "uniform on the simplex (Dirichlet(1)), as ToR-OT's training data",
+  "official_config": {
+    "seed": 0,
+    "data": "random",
+    "potential_model": {
+      "_target_": "meta_ot.models.PotentialMLP",
+      "n_hidden": 1024,
+      "n_hidden_layer": 3
+    },
+    "batch_size": 128,
+    "num_train_iter": 50000,
+    "potential_model_mlp": {
+      "_target_": "meta_ot.models.PotentialMLP",
+      "n_hidden": 1024,
+      "n_hidden_layer": 3
+    },
+    "optim": {
+      "_target_": "optax.adam",
+      "learning_rate": 0.001
+    },
+    "max_grad_norm": 1.0
+  },
+  "vendored_commit": "7dcc9a4c25b0a57626d3fa26028c6633ad19da50",
+  "versions": {
+    "jax": "0.4.30",
+    "jaxlib": "0.4.30",
+    "flax": "0.8.5",
+    "optax": "0.2.2",
+    "numpy": "1.26.4",
+    "ott-jax": "0.2.9"
+  },
+  "device": "cuda:0",
+  "provenance": {
+    "tor_ot_commit": "b9c045807b7fbee37ab22a523e75bc3bbc076169",
+    "tor_ot_dirty": true,
+    "python": "3.10.22",
+    "executable": "/home/ivan/anaconda3/envs/env_metaot/bin/python",
+    "platform": "Linux-6.8.0-138-generic-x86_64-with-glibc2.35"
+  }
+}

baselines/metaot/cubic_euclidean_eps0.001/metrics.json

@@ -0,0 +1,15 @@
+{
+  "validation": {
+    "set": "val_simplex2d_cubic_euclidean",
+    "instances": 500,
+    "paper_mape_mean": 0.43835759711989714,
+    "paper_mape_median": 0.42727180654938646,
+    "row_violation_l1_mean": 0.30611072450875987,
+    "col_violation_l1_mean": 1.227751313463666e-15,
+    "abs_cost_gap_mean": 0.0006078786352234964,
+    "finite_fraction": 1.0
+  },
+  "train_seconds": 4130.424247264862,
+  "parameters": 2210852,
+  "iterations": 50000
+}

baselines/metaot/cubic_euclidean_eps0.001/params.msgpack

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dabb07a7b2f614b3c8dc14fcdd5a51b44d8820ac04fdee0087908be5b06460d9
+size 8843673

baselines/metaot/euclidean_eps0.01/config.json

@@ -0,0 +1,47 @@
+{
+  "method": "metaot",
+  "cost": "euclidean",
+  "eps": 0.01,
+  "n": 36,
+  "dim": 2,
+  "seed": 0,
+  "train_marginals": "uniform on the simplex (Dirichlet(1)), as ToR-OT's training data",
+  "official_config": {
+    "seed": 0,
+    "data": "random",
+    "potential_model": {
+      "_target_": "meta_ot.models.PotentialMLP",
+      "n_hidden": 1024,
+      "n_hidden_layer": 3
+    },
+    "batch_size": 128,
+    "num_train_iter": 50000,
+    "potential_model_mlp": {
+      "_target_": "meta_ot.models.PotentialMLP",
+      "n_hidden": 1024,
+      "n_hidden_layer": 3
+    },
+    "optim": {
+      "_target_": "optax.adam",
+      "learning_rate": 0.001
+    },
+    "max_grad_norm": 1.0
+  },
+  "vendored_commit": "7dcc9a4c25b0a57626d3fa26028c6633ad19da50",
+  "versions": {
+    "jax": "0.4.30",
+    "jaxlib": "0.4.30",
+    "flax": "0.8.5",
+    "optax": "0.2.2",
+    "numpy": "1.26.4",
+    "ott-jax": "0.2.9"
+  },
+  "device": "cuda:0",
+  "provenance": {
+    "tor_ot_commit": "b9c045807b7fbee37ab22a523e75bc3bbc076169",
+    "tor_ot_dirty": true,
+    "python": "3.10.22",
+    "executable": "/home/ivan/anaconda3/envs/env_metaot/bin/python",
+    "platform": "Linux-6.8.0-138-generic-x86_64-with-glibc2.35"
+  }
+}

baselines/metaot/euclidean_eps0.01/metrics.json

@@ -0,0 +1,15 @@
+{
+  "validation": {
+    "set": "val_simplex2d_euclidean",
+    "instances": 500,
+    "paper_mape_mean": 0.41057894637272296,
+    "paper_mape_median": 0.3993071094091103,
+    "row_violation_l1_mean": 0.24206822685050514,
+    "col_violation_l1_mean": 1.1271179963071335e-15,
+    "abs_cost_gap_mean": 0.008357788421720032,
+    "finite_fraction": 1.0
+  },
+  "train_seconds": 4130.422701120377,
+  "parameters": 2210852,
+  "iterations": 50000
+}

baselines/metaot/euclidean_eps0.01/params.msgpack

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:90ed21e379ece73534f5eb030bd17fcb4f80711aa6ce04abde490cc3617e67bc
+size 8843673

baselines/metaot/gaussian_eps0.01/config.json

@@ -0,0 +1,47 @@
+{
+  "method": "metaot",
+  "cost": "gaussian",
+  "eps": 0.01,
+  "n": 36,
+  "dim": 2,
+  "seed": 0,
+  "train_marginals": "uniform on the simplex (Dirichlet(1)), as ToR-OT's training data",
+  "official_config": {
+    "seed": 0,
+    "data": "random",
+    "potential_model": {
+      "_target_": "meta_ot.models.PotentialMLP",
+      "n_hidden": 1024,
+      "n_hidden_layer": 3
+    },
+    "batch_size": 128,
+    "num_train_iter": 50000,
+    "potential_model_mlp": {
+      "_target_": "meta_ot.models.PotentialMLP",
+      "n_hidden": 1024,
+      "n_hidden_layer": 3
+    },
+    "optim": {
+      "_target_": "optax.adam",
+      "learning_rate": 0.001
+    },
+    "max_grad_norm": 1.0
+  },
+  "vendored_commit": "7dcc9a4c25b0a57626d3fa26028c6633ad19da50",
+  "versions": {
+    "jax": "0.4.30",
+    "jaxlib": "0.4.30",
+    "flax": "0.8.5",
+    "optax": "0.2.2",
+    "numpy": "1.26.4",
+    "ott-jax": "0.2.9"
+  },
+  "device": "cuda:0",
+  "provenance": {
+    "tor_ot_commit": "b9c045807b7fbee37ab22a523e75bc3bbc076169",
+    "tor_ot_dirty": true,
+    "python": "3.10.22",
+    "executable": "/home/ivan/anaconda3/envs/env_metaot/bin/python",
+    "platform": "Linux-6.8.0-138-generic-x86_64-with-glibc2.35"
+  }
+}

baselines/metaot/gaussian_eps0.01/metrics.json

@@ -0,0 +1,15 @@
+{
+  "validation": {
+    "set": "val_simplex2d_gaussian",
+    "instances": 500,
+    "paper_mape_mean": 0.5226949643986197,
+    "paper_mape_median": 0.5212648175950179,
+    "row_violation_l1_mean": 0.4143645651665812,
+    "col_violation_l1_mean": 2.194984908011877e-15,
+    "abs_cost_gap_mean": 0.030263616592299654,
+    "finite_fraction": 1.0
+  },
+  "train_seconds": 4130.375089645386,
+  "parameters": 2210852,
+  "iterations": 50000
+}

baselines/metaot/gaussian_eps0.01/params.msgpack

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:03998a558012f61ca2513f809d9f4f116fe33370440ab4600509a33362e81b0a
+size 8843673

baselines/metaot/sqeuclidean_eps0.01/config.json

@@ -0,0 +1,47 @@
+{
+  "method": "metaot",
+  "cost": "sqeuclidean",
+  "eps": 0.01,
+  "n": 36,
+  "dim": 2,
+  "seed": 0,
+  "train_marginals": "uniform on the simplex (Dirichlet(1)), as ToR-OT's training data",
+  "official_config": {
+    "seed": 0,
+    "data": "random",
+    "potential_model": {
+      "_target_": "meta_ot.models.PotentialMLP",
+      "n_hidden": 1024,
+      "n_hidden_layer": 3
+    },
+    "batch_size": 128,
+    "num_train_iter": 50000,
+    "potential_model_mlp": {
+      "_target_": "meta_ot.models.PotentialMLP",
+      "n_hidden": 1024,
+      "n_hidden_layer": 3
+    },
+    "optim": {
+      "_target_": "optax.adam",
+      "learning_rate": 0.001
+    },
+    "max_grad_norm": 1.0
+  },
+  "vendored_commit": "7dcc9a4c25b0a57626d3fa26028c6633ad19da50",
+  "versions": {
+    "jax": "0.4.30",
+    "jaxlib": "0.4.30",
+    "flax": "0.8.5",
+    "optax": "0.2.2",
+    "numpy": "1.26.4",
+    "ott-jax": "0.2.9"
+  },
+  "device": "cuda:0",
+  "provenance": {
+    "tor_ot_commit": "b9c045807b7fbee37ab22a523e75bc3bbc076169",
+    "tor_ot_dirty": true,
+    "python": "3.10.22",
+    "executable": "/home/ivan/anaconda3/envs/env_metaot/bin/python",
+    "platform": "Linux-6.8.0-138-generic-x86_64-with-glibc2.35"
+  }
+}

baselines/metaot/sqeuclidean_eps0.01/metrics.json

@@ -0,0 +1,15 @@
+{
+  "validation": {
+    "set": "val_simplex2d_sqeuclidean",
+    "instances": 500,
+    "paper_mape_mean": 0.49454423258223007,
+    "paper_mape_median": 0.4890397708741878,
+    "row_violation_l1_mean": 0.10032652592309427,
+    "col_violation_l1_mean": 4.806512341288475e-16,
+    "abs_cost_gap_mean": 0.002031926991414964,
+    "finite_fraction": 1.0
+  },
+  "train_seconds": 4130.550829410553,
+  "parameters": 2210852,
+  "iterations": 50000
+}

baselines/metaot/sqeuclidean_eps0.01/params.msgpack

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9cbf8a93696f96dcbd934d257d97f3ccbd3a22f40ec8b4dfe11fb1e4573e3872
+size 8843673

baselines/metaot/sqrt_euclidean_eps0.01/config.json

@@ -0,0 +1,47 @@
+{
+  "method": "metaot",
+  "cost": "sqrt_euclidean",
+  "eps": 0.01,
+  "n": 36,
+  "dim": 2,
+  "seed": 0,
+  "train_marginals": "uniform on the simplex (Dirichlet(1)), as ToR-OT's training data",
+  "official_config": {
+    "seed": 0,
+    "data": "random",
+    "potential_model": {
+      "_target_": "meta_ot.models.PotentialMLP",
+      "n_hidden": 1024,
+      "n_hidden_layer": 3
+    },
+    "batch_size": 128,
+    "num_train_iter": 50000,
+    "potential_model_mlp": {
+      "_target_": "meta_ot.models.PotentialMLP",
+      "n_hidden": 1024,
+      "n_hidden_layer": 3
+    },
+    "optim": {
+      "_target_": "optax.adam",
+      "learning_rate": 0.001
+    },
+    "max_grad_norm": 1.0
+  },
+  "vendored_commit": "7dcc9a4c25b0a57626d3fa26028c6633ad19da50",
+  "versions": {
+    "jax": "0.4.30",
+    "jaxlib": "0.4.30",
+    "flax": "0.8.5",
+    "optax": "0.2.2",
+    "numpy": "1.26.4",
+    "ott-jax": "0.2.9"
+  },
+  "device": "cuda:0",
+  "provenance": {
+    "tor_ot_commit": "b9c045807b7fbee37ab22a523e75bc3bbc076169",
+    "tor_ot_dirty": true,
+    "python": "3.10.22",
+    "executable": "/home/ivan/anaconda3/envs/env_metaot/bin/python",
+    "platform": "Linux-6.8.0-138-generic-x86_64-with-glibc2.35"
+  }
+}

baselines/metaot/sqrt_euclidean_eps0.01/metrics.json

@@ -0,0 +1,15 @@
+{
+  "validation": {
+    "set": "val_simplex2d_sqrt_euclidean",
+    "instances": 500,
+    "paper_mape_mean": 0.360679581856843,
+    "paper_mape_median": 0.35895149743904375,
+    "row_violation_l1_mean": 0.27511418609134025,
+    "col_violation_l1_mean": 2.0080832422471187e-15,
+    "abs_cost_gap_mean": 0.01639749865033589,
+    "finite_fraction": 1.0
+  },
+  "train_seconds": 4130.353412151337,
+  "parameters": 2210852,
+  "iterations": 50000
+}

baselines/metaot/sqrt_euclidean_eps0.01/params.msgpack

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:12de5d3b6b2c5a21bf3f82d644697011b99d71f4254d26e64ca2872ac8f98205
+size 8843673

baselines/unot/canberra_eps0.01/config.json

@@ -0,0 +1,41 @@
+{
+  "method": "unot",
+  "cost": "canberra",
+  "eps": 0.01,
+  "n": 36,
+  "dim": 2,
+  "seed": 0,
+  "hyperparameters": {
+    "lr": 0.01,
+    "weight_decay_generator": 0.0,
+    "weight_decay_predictor": 0.0001,
+    "gamma_generator": 1.0,
+    "gamma_predictor": 0.9999,
+    "batch_size": 64,
+    "nits": 10000,
+    "nits_mini_loop_generator": 1,
+    "nits_mini_loop_predictor": 1,
+    "sinkhorn_max_iterations": 5,
+    "num_layers": 3,
+    "length_latent": 10,
+    "dust_const": 1e-06,
+    "skip_const": 1.0,
+    "width": 1024
+  },
+  "train_marginals": "UNOT's adversarial generator (no dataset), as in the official method",
+  "sinkhorn_targets": "log domain (identical iterates)",
+  "vendored_commit": "73b290ba8c0709fa7dd4992eeda88171628d6bb2",
+  "versions": {
+    "torch": "2.5.1+cu124",
+    "numpy": "2.2.1",
+    "wandb": "0.19.1"
+  },
+  "device": "NVIDIA A100-PCIE-40GB",
+  "provenance": {
+    "tor_ot_commit": "b9c045807b7fbee37ab22a523e75bc3bbc076169",
+    "tor_ot_dirty": true,
+    "python": "3.12.15",
+    "executable": "/home/ivan/anaconda3/envs/env_unot/bin/python",
+    "platform": "Linux-6.8.0-138-generic-x86_64-with-glibc2.35"
+  }
+}

baselines/unot/canberra_eps0.01/generator.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e3d60a959a1da9310f9d4c5f4df82b55098fdf8be48d82babad12be06375d032
+size 13789312

baselines/unot/canberra_eps0.01/metrics.json

@@ -0,0 +1,19 @@
+{
+  "validation": {
+    "set": "val_simplex2d_canberra",
+    "instances": 500,
+    "paper_mape_mean": 1.0239031142795592,
+    "paper_mape_median": 1.0110007108679273,
+    "row_violation_l1_mean": 2.1793333430843205e-08,
+    "col_violation_l1_mean": 0.8667472874915544,
+    "abs_cost_gap_mean": 0.032092018875914785,
+    "finite_fraction": 1.0
+  },
+  "train_seconds": 1100.571202993393,
+  "parameters": 3268644,
+  "iterations": 10000,
+  "predictor_updates": 10000,
+  "non_finite_predictor_losses": 0,
+  "final_predictor_loss_mean_last_100": 12.474288320541381,
+  "batchnorm_recalibrated_on": "12800 finite generator samples (0 non-finite dropped)"
+}

baselines/unot/canberra_eps0.01/predictor.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0b4fcbea2e5ba01bc960a5fea9c9dacee9588da2228de3dc8e5b705b790cf07e
+size 13117440

baselines/unot/cubic_euclidean_eps0.001/config.json

@@ -0,0 +1,41 @@
+{
+  "method": "unot",
+  "cost": "cubic_euclidean",
+  "eps": 0.001,
+  "n": 36,
+  "dim": 2,
+  "seed": 0,
+  "hyperparameters": {
+    "lr": 0.01,
+    "weight_decay_generator": 0.0,
+    "weight_decay_predictor": 0.0001,
+    "gamma_generator": 1.0,
+    "gamma_predictor": 0.9999,
+    "batch_size": 64,
+    "nits": 10000,
+    "nits_mini_loop_generator": 1,
+    "nits_mini_loop_predictor": 1,
+    "sinkhorn_max_iterations": 5,
+    "num_layers": 3,
+    "length_latent": 10,
+    "dust_const": 1e-06,
+    "skip_const": 1.0,
+    "width": 1024
+  },
+  "train_marginals": "UNOT's adversarial generator (no dataset), as in the official method",
+  "sinkhorn_targets": "log domain (identical iterates)",
+  "vendored_commit": "73b290ba8c0709fa7dd4992eeda88171628d6bb2",
+  "versions": {
+    "torch": "2.5.1+cu124",
+    "numpy": "2.2.1",
+    "wandb": "0.19.1"
+  },
+  "device": "NVIDIA A100-PCIE-40GB",
+  "provenance": {
+    "tor_ot_commit": "b9c045807b7fbee37ab22a523e75bc3bbc076169",
+    "tor_ot_dirty": true,
+    "python": "3.12.15",
+    "executable": "/home/ivan/anaconda3/envs/env_unot/bin/python",
+    "platform": "Linux-6.8.0-138-generic-x86_64-with-glibc2.35"
+  }
+}

baselines/unot/cubic_euclidean_eps0.001/generator.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2a189fe332ab67c3c9ef4fe177f6d21491c0d6bb39c4906513e32af4a8f66bea
+size 13789312

baselines/unot/cubic_euclidean_eps0.001/metrics.json

@@ -0,0 +1,19 @@
+{
+  "validation": {
+    "set": "val_simplex2d_cubic_euclidean",
+    "instances": 500,
+    "paper_mape_mean": 0.9532561271039957,
+    "paper_mape_median": 0.9460185019162484,
+    "row_violation_l1_mean": 2.1793333400397737e-08,
+    "col_violation_l1_mean": 0.762428565899103,
+    "abs_cost_gap_mean": 0.0012899823837028265,
+    "finite_fraction": 1.0
+  },
+  "train_seconds": 1100.405916929245,
+  "parameters": 3268644,
+  "iterations": 10000,
+  "predictor_updates": 10000,
+  "non_finite_predictor_losses": 0,
+  "final_predictor_loss_mean_last_100": 17.718043403625487,
+  "batchnorm_recalibrated_on": "12800 finite generator samples (0 non-finite dropped)"
+}

baselines/unot/cubic_euclidean_eps0.001/predictor.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3c974ffc49bf17f273f474aadd6a7634214261b2896d5516ce33ed32c65dc1bd
+size 13117440

baselines/unot/euclidean_eps0.01/config.json

@@ -0,0 +1,41 @@
+{
+  "method": "unot",
+  "cost": "euclidean",
+  "eps": 0.01,
+  "n": 36,
+  "dim": 2,
+  "seed": 0,
+  "hyperparameters": {
+    "lr": 0.01,
+    "weight_decay_generator": 0.0,
+    "weight_decay_predictor": 0.0001,
+    "gamma_generator": 1.0,
+    "gamma_predictor": 0.9999,
+    "batch_size": 64,
+    "nits": 10000,
+    "nits_mini_loop_generator": 1,
+    "nits_mini_loop_predictor": 1,
+    "sinkhorn_max_iterations": 5,
+    "num_layers": 3,
+    "length_latent": 10,
+    "dust_const": 1e-06,
+    "skip_const": 1.0,
+    "width": 1024
+  },
+  "train_marginals": "UNOT's adversarial generator (no dataset), as in the official method",
+  "sinkhorn_targets": "log domain (identical iterates)",
+  "vendored_commit": "73b290ba8c0709fa7dd4992eeda88171628d6bb2",
+  "versions": {
+    "torch": "2.5.1+cu124",
+    "numpy": "2.2.1",
+    "wandb": "0.19.1"
+  },
+  "device": "NVIDIA A100-PCIE-40GB",
+  "provenance": {
+    "tor_ot_commit": "b9c045807b7fbee37ab22a523e75bc3bbc076169",
+    "tor_ot_dirty": true,
+    "python": "3.12.15",
+    "executable": "/home/ivan/anaconda3/envs/env_unot/bin/python",
+    "platform": "Linux-6.8.0-138-generic-x86_64-with-glibc2.35"
+  }
+}

baselines/unot/euclidean_eps0.01/generator.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:beb27a049f735baf49f85aebe9e4734941595dbbc301ac2db3e89cd46a396130
+size 13789312

baselines/unot/euclidean_eps0.01/metrics.json

@@ -0,0 +1,19 @@
+{
+  "validation": {
+    "set": "val_simplex2d_euclidean",
+    "instances": 500,
+    "paper_mape_mean": 0.7710163460437718,
+    "paper_mape_median": 0.7651310466306106,
+    "row_violation_l1_mean": 2.1793333438916048e-08,
+    "col_violation_l1_mean": 0.6470750562478266,
+    "abs_cost_gap_mean": 0.03352212004773903,
+    "finite_fraction": 1.0
+  },
+  "train_seconds": 1100.8721661567688,
+  "parameters": 3268644,
+  "iterations": 10000,
+  "predictor_updates": 10000,
+  "non_finite_predictor_losses": 0,
+  "final_predictor_loss_mean_last_100": 12.550443811416626,
+  "batchnorm_recalibrated_on": "12800 finite generator samples (0 non-finite dropped)"
+}

baselines/unot/euclidean_eps0.01/predictor.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e38287705ccb366e2682f264ef1023633ece05e5e319ffe6b0df90084b102223
+size 13117440

baselines/unot/gaussian_eps0.01_w144/config.json

@@ -0,0 +1,41 @@
+{
+  "method": "unot",
+  "cost": "gaussian",
+  "eps": 0.01,
+  "n": 36,
+  "dim": 2,
+  "seed": 0,
+  "hyperparameters": {
+    "lr": 0.01,
+    "weight_decay_generator": 0.0,
+    "weight_decay_predictor": 0.0001,
+    "gamma_generator": 1.0,
+    "gamma_predictor": 0.9999,
+    "batch_size": 64,
+    "nits": 10000,
+    "nits_mini_loop_generator": 1,
+    "nits_mini_loop_predictor": 1,
+    "sinkhorn_max_iterations": 5,
+    "num_layers": 3,
+    "length_latent": 10,
+    "dust_const": 1e-06,
+    "skip_const": 1.0,
+    "width": 144
+  },
+  "train_marginals": "UNOT's adversarial generator (no dataset), as in the official method",
+  "sinkhorn_targets": "log domain (identical iterates)",
+  "vendored_commit": "73b290ba8c0709fa7dd4992eeda88171628d6bb2",
+  "versions": {
+    "torch": "2.5.1+cu124",
+    "numpy": "2.2.1",
+    "wandb": "0.19.1"
+  },
+  "device": "NVIDIA A100-PCIE-40GB",
+  "provenance": {
+    "tor_ot_commit": "b9c045807b7fbee37ab22a523e75bc3bbc076169",
+    "tor_ot_dirty": true,
+    "python": "3.12.15",
+    "executable": "/home/ivan/anaconda3/envs/env_unot/bin/python",
+    "platform": "Linux-6.8.0-138-generic-x86_64-with-glibc2.35"
+  }
+}

baselines/unot/gaussian_eps0.01_w144/generator.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d1ac19634c248e2daef3c212ff7c4984868a2d2e47fa9500732d45326a3df219
+size 427328

baselines/unot/gaussian_eps0.01_w144/metrics.json

@@ -0,0 +1,19 @@
+{
+  "validation": {
+    "set": "val_simplex2d_gaussian",
+    "instances": 500,
+    "paper_mape_mean": 0.8667829143279226,
+    "paper_mape_median": 0.8609271789295385,
+    "row_violation_l1_mean": 2.1793333412383664e-08,
+    "col_violation_l1_mean": 0.7243300167267797,
+    "abs_cost_gap_mean": 0.09239843748887588,
+    "finite_fraction": 1.0
+  },
+  "train_seconds": 659.0352730751038,
+  "parameters": 79524,
+  "iterations": 10000,
+  "predictor_updates": 10000,
+  "non_finite_predictor_losses": 0,
+  "final_predictor_loss_mean_last_100": 30.131474838256835,
+  "batchnorm_recalibrated_on": "12800 finite generator samples (0 non-finite dropped)"
+}

baselines/unot/gaussian_eps0.01_w144/predictor.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:daeb29419504918302b5623684e9dd3645239e6dfcfa2155f3bdb721ca18fa4b
+size 332736

baselines/unot/sqeuclidean_eps0.01_w144/config.json

@@ -0,0 +1,41 @@
+{
+  "method": "unot",
+  "cost": "sqeuclidean",
+  "eps": 0.01,
+  "n": 36,
+  "dim": 2,
+  "seed": 0,
+  "hyperparameters": {
+    "lr": 0.01,
+    "weight_decay_generator": 0.0,
+    "weight_decay_predictor": 0.0001,
+    "gamma_generator": 1.0,
+    "gamma_predictor": 0.9999,
+    "batch_size": 64,
+    "nits": 10000,
+    "nits_mini_loop_generator": 1,
+    "nits_mini_loop_predictor": 1,
+    "sinkhorn_max_iterations": 5,
+    "num_layers": 3,
+    "length_latent": 10,
+    "dust_const": 1e-06,
+    "skip_const": 1.0,
+    "width": 144
+  },
+  "train_marginals": "UNOT's adversarial generator (no dataset), as in the official method",
+  "sinkhorn_targets": "log domain (identical iterates)",
+  "vendored_commit": "73b290ba8c0709fa7dd4992eeda88171628d6bb2",
+  "versions": {
+    "torch": "2.5.1+cu124",
+    "numpy": "2.2.1",
+    "wandb": "0.19.1"
+  },
+  "device": "NVIDIA A100-PCIE-40GB",
+  "provenance": {
+    "tor_ot_commit": "b9c045807b7fbee37ab22a523e75bc3bbc076169",
+    "tor_ot_dirty": true,
+    "python": "3.12.15",
+    "executable": "/home/ivan/anaconda3/envs/env_unot/bin/python",
+    "platform": "Linux-6.8.0-138-generic-x86_64-with-glibc2.35"
+  }
+}

baselines/unot/sqeuclidean_eps0.01_w144/generator.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9fc5bbd5c191063a7f5ee247c2ca499b58cebcf466d5bba2cf6af06d5161624a
+size 427328

baselines/unot/sqeuclidean_eps0.01_w144/metrics.json

@@ -0,0 +1,19 @@
+{
+  "validation": {
+    "set": "val_simplex2d_sqeuclidean",
+    "instances": 500,
+    "paper_mape_mean": 0.9227674866907752,
+    "paper_mape_median": 0.9187091289427531,
+    "row_violation_l1_mean": 2.1793333416468553e-08,
+    "col_violation_l1_mean": 0.5871081061916796,
+    "abs_cost_gap_mean": 0.008746131607150119,
+    "finite_fraction": 1.0
+  },
+  "train_seconds": 658.9568598270416,
+  "parameters": 79524,
+  "iterations": 10000,
+  "predictor_updates": 10000,
+  "non_finite_predictor_losses": 0,
+  "final_predictor_loss_mean_last_100": 12.587937898635865,
+  "batchnorm_recalibrated_on": "12800 finite generator samples (0 non-finite dropped)"
+}

baselines/unot/sqeuclidean_eps0.01_w144/predictor.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9788fcdfff26b79823071cdb3e75d5c8bcab98cca86d9d575152d198db3bcedb
+size 332736

baselines/unot/sqrt_euclidean_eps0.01/config.json

@@ -0,0 +1,41 @@
+{
+  "method": "unot",
+  "cost": "sqrt_euclidean",
+  "eps": 0.01,
+  "n": 36,
+  "dim": 2,
+  "seed": 0,
+  "hyperparameters": {
+    "lr": 0.01,
+    "weight_decay_generator": 0.0,
+    "weight_decay_predictor": 0.0001,
+    "gamma_generator": 1.0,
+    "gamma_predictor": 0.9999,
+    "batch_size": 64,
+    "nits": 10000,
+    "nits_mini_loop_generator": 1,
+    "nits_mini_loop_predictor": 1,
+    "sinkhorn_max_iterations": 5,
+    "num_layers": 3,
+    "length_latent": 10,
+    "dust_const": 1e-06,
+    "skip_const": 1.0,
+    "width": 1024
+  },
+  "train_marginals": "UNOT's adversarial generator (no dataset), as in the official method",
+  "sinkhorn_targets": "log domain (identical iterates)",
+  "vendored_commit": "73b290ba8c0709fa7dd4992eeda88171628d6bb2",
+  "versions": {
+    "torch": "2.5.1+cu124",
+    "numpy": "2.2.1",
+    "wandb": "0.19.1"
+  },
+  "device": "NVIDIA A100-PCIE-40GB",
+  "provenance": {
+    "tor_ot_commit": "b9c045807b7fbee37ab22a523e75bc3bbc076169",
+    "tor_ot_dirty": true,
+    "python": "3.12.15",
+    "executable": "/home/ivan/anaconda3/envs/env_unot/bin/python",
+    "platform": "Linux-6.8.0-138-generic-x86_64-with-glibc2.35"
+  }
+}

baselines/unot/sqrt_euclidean_eps0.01/generator.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c031c18f0ec18c52b070eb324a85d845148edaad7f21eb0bb91385ce19f8fe1
+size 13789312

baselines/unot/sqrt_euclidean_eps0.01/metrics.json

@@ -0,0 +1,19 @@
+{
+  "validation": {
+    "set": "val_simplex2d_sqrt_euclidean",
+    "instances": 500,
+    "paper_mape_mean": 0.8049276410675595,
+    "paper_mape_median": 0.8038773890528728,
+    "row_violation_l1_mean": 2.179333341410435e-08,
+    "col_violation_l1_mean": 0.731193874708468,
+    "abs_cost_gap_mean": 0.0926930243704554,
+    "finite_fraction": 1.0
+  },
+  "train_seconds": 1100.6754524707794,
+  "parameters": 3268644,
+  "iterations": 10000,
+  "predictor_updates": 10000,
+  "non_finite_predictor_losses": 0,
+  "final_predictor_loss_mean_last_100": 21.89601837158203,
+  "batchnorm_recalibrated_on": "12800 finite generator samples (0 non-finite dropped)"
+}

baselines/unot/sqrt_euclidean_eps0.01/predictor.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b2ec80b1427c1f07f99ed361688067b1304c5241cdd7ce470038087007a27b75
+size 13117440

config.yaml

@@ -0,0 +1,56 @@
+data:
+  repo_id: ivanjaenm/torot_data_4M_2d
+  cache_dir: null
+  train_split: train
+  val_split: val
+  max_bins: 200
+  transpose_aug_prob: 0.5
+  residual_noise: 0.01
+  residual_leftover_prob: 0.25
+  residual_leftover_range:
+  - 1.0e-06
+  - 0.01
+  num_workers: 8
+tokenizer:
+  mantissa_digits: 4
+  max_exponent: 9
+model:
+  d_model: 768
+  n_heads: 12
+  enc_layers: 6
+  dec_layers: 12
+  d_ff: 2048
+  fourier_bands: 8
+  dropout: 0.0
+  max_seq_len: 4096
+  decoding: row
+loss:
+  digit_weight_decay: 0.5
+  exponent_weight: 1.0
+  ordinal_eps: 0.05
+  z_loss: 0.0001
+optim:
+  lr: 0.0001
+  betas:
+  - 0.9
+  - 0.95
+  weight_decay: 0.1
+  warmup_steps: 2000
+  max_steps: 627000
+  min_lr_ratio: 0.1
+  batch_size: 256
+  max_tokens: 32768
+  max_dec_tokens: 24576
+  grad_clip: 1.0
+  ema_decay: 0.999
+run:
+  out_dir: experiments/torot/row_M_2d_mixed_support
+  seed: 0
+  log_every: 50
+  val_every: 5000
+  val_batches: 20
+  val_decode_instances: 256
+  max_decode_len: 2048
+  bf16: true
+  compile: false
+  resume: null

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:99bfdb2eedf149f2be416a4f69ad4ce37ed6b64a252c474b2cf95191fd6c92f5
+size 630429520