Toolbox for Gromov-Wasserstein optimal transport: Application to unsupervised alignment in neuroscience

Abstract

AbstractGromov-Wasserstein optimal transport (GWOT) has emerged as a versatile method for unsupervised alignment in various research areas, including neuroscience, drawing upon the strengths of optimal transport theory. However, the use of GWOT in various applications has been hindered by the difficulty of finding good optima, a significant challenge stemming from GWOT’s nature as a non-convex optimization method. It is often difficult to avoid suboptimal local optima because of the need for systematic hyperparameter tuning. To overcome these obstacles, this paper presents a user-friendly GWOT hyperparameter tuning toolbox (GWTune) specifically designed to streamline the use of GWOT in neuroscience and other fields. The toolbox incorporates Optuna, an advanced hyperparameter tuning tool that uses Bayesian sampling to increase the chances of finding favorable local optima. To demonstrate the utility of our toolbox, we first illustrate the qualitative difference between the conventional supervised alignment method and our unsupervised alignment method using synthetic data. Then, we demonstrate the applicability of our toolbox using some typical examples in neuroscience. Specifically, we applied GWOT to the similarity structures of natural objects or natural scenes obtained from three data domains: behavioral data, neural data, and neural network models. This toolbox is an accessible and robust solution for practical applications in neuroscience and beyond, making the powerful GWOT methodology more accessible to a wider range of users. The open source code for the toolbox is available on GitHub. This work not only facilitates the application of GWOT, but also opens avenues for future improvements and extensions.