ToR-OT 2D

ToR-OT 2D — исследовательская модель на 157.6M параметров из статьи Learning Optimal Transport Plans Via Autoregressive Token Regression (NeurIPS 2026). Encoder читает две гистограммы и тип cost, а decoder выдаёт optimal transport plan построчно в виде токенов: указатель на target-bin, экспонента и четыре мантиссы. Модель предназначена для 2D-результатов статьи.

✨ Ключевые особенности

• 📓 Статья: Learning Optimal Transport Plans Via Autoregressive Token Regression (NeurIPS 2026).
• 🤖 Encoder принимает две гистограммы и тип cost; decoder пишет план построчно токенами (target-bin pointer, exponent, 4 mantissa digits).
• 🎲 Обучение: 2D-гистограммы на k×k сетках в [0,1]^2, k=3–10, n=9–100 bins, шесть cost types; в половине маргиналов 10–70% bins пустые.
• ▶ Декодирование в статье: greedy, grammar-constrained, row by row; последняя строка — моделью (last_row='model').
• 💻 Веса float32, EMA после 620000 шагов; в репозитории есть код загрузки/предсказания и baselines Meta OT и UNOT.

💻 Пример

import torch
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)
plan = predict_plans(model, tok, [inst], device, last_row='model')[0]

Подробнее

Бенчмарков в карточке нет; пример кода считает relative cost gap к точному плану.
Веса float32, 157.6M параметров; требования к VRAM не указаны.
Обучение: EMA после 620000 шагов; данные — 2D-гистограммы k=3–10, n=9–100, шесть cost types.
Ограничение: модель для 2D-планов из статьи, не универсальный OT-солвер.