最优传输工具 (OTT):一个用于各类 Wasserstein 计算的 JAX 工具箱
机器学习
2022-01-31 v1 机器学习
摘要
最优传输工具 (OTT-JAX) 是一个 Python 工具箱,可求解点云与直方图之间的最优传输问题。该工具箱构建于 JAX 的多种特性之上,例如自动与自定义反向模式微分、向量化、即时编译以及对加速器的支持。该工具箱涵盖基础计算(如正则化最优传输问题的求解)以及更高级的扩展,如重心、Gromov-Wasserstein、低秩求解器、凸映射估计、分位数与秩的可微推广,以及高斯混合之间的近似最优传输。工具箱代码可在 \texttt{https://github.com/ott-jax/ott} 获取。
引用
@article{arxiv.2201.12324,
title = {Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein},
author = {Marco Cuturi and Laetitia Meng-Papaxanthos and Yingtao Tian and Charlotte Bunne and Geoff Davis and Olivier Teboul},
journal= {arXiv preprint arXiv:2201.12324},
year = {2022}
}
备注
4 pages