Прорыв в масштабируемости оптимального транспорта
Классические решатели энтропийного оптимального транспорта (EOT), такие как оригинальный FlashSinkhorn, требуют вычисления всех $n \times m$ пар точек на каждой итерации, что ограничивает их применение на больших наборах данных. Новая версия FlashSinkhorn 2 (FS2) предлагает двухэтапный подход для задач с евклидовым расстоянием, позволяющий обрабатывать облака точек низкой размерности на одной GPU без хранения плотной матрицы ядра.
Двухэтапная архитектура: от грубой оценки к точности
Алгоритм FS2 сочетает два ключевых этапа для оптимизации вычислений:
- Грубый этап (Coarse stage): Решение строится на центроидах ячеек. Потенциалы «поднимаются» (lift) ко всем точкам. Если проверка маргиналов отклоняет это приближение, алгоритм продолжает работу на центроидах, избегая большинства обновлений на уровне отдельных точек.
- Точный этап (Fine stage): Использует блочно-разреженную структуру (block-sparse). Блоки, упорядоченные по пространственной кривой Мортона, позволяют применять экранирование (screening) и объединенное выполнение через тензорные ядра (fused tensor-core execution). Порог, заданный массами блоков, ограничивает вклад каждого опущенного тайла в каждую строку и столбец.
Рекордные метрики на GPU A100
На синтетических бенчмарках FS2 достиг целевого остатка во всех 32 задачах и в 10 задачах мультимасштабного GeomLoss. Главным достижением стало решение дискретной задачи EOT для мер космической N-тел симуляции с 134 миллионами частиц ($1.34 \times 10^8$). При энтропийном размытии, равном среднему межчастичному расстоянию, алгоритм достиг остатка маргиналов менее 0.01 за менее 2.5 часов на одной GPU NVIDIA A100.
| Параметр | Значение / Характеристика |
|---|---|
| Название модели | FlashSinkhorn 2 (FS2) |
| Тип задачи | Дискретный EOT (Squared-Euclidean cost) |
| Объем данных (рекорд) | 134 млн частиц ($1.34 \times 10^8$) |
| Аппаратная платформа | Single NVIDIA A100 |
| Время решения | < 2.5 часов |
| Точность (residual) | < 0.01 (all-particle marginal residual) |
| Ключевая оптимизация | Block-sparse, Morton-ordering, Tensor Cores |
Значение для науки и индустрии
По словам авторов, это, вероятно, самая большая дискретная задача EOT, решенная с такой точностью за несколько часов. Открытая реализация алгоритма доступна для воспроизведения результатов, что открывает новые возможности для космологии, обработки больших данных и машинного обучения, где требуется эффективное сопоставление огромных наборов точек.
Источник: arXiv cs.AI ↗