Wasserstein barycenters provide a principled approach for aggregating probability measures, while preserving the geometry of their ambient space. Existing discrete methods are not scalable as they assume access to the complete set of samples from the input measures. Meanwhile, neural network approaches do scale well, but rely on complex optimization problems and cannot easily incorporate label information. We address these limitations through gradient flows in the space of probability measures. Through time discretization, we achieve a scalable algorithm that i) relies on mini-batch optimal transport, ii) accepts modular regularization through task-aware functions, and iii) seamlessly integrates supervised information into the ground-cost. We empirically validate our approach on domain adaptation benchmarks that span computer vision, neuroscience, and chemical engineering. Our method establishes a new state-of-the-art barycenter solver, with labeled barycenters consistently outperforming unlabeled ones.

翻译：Wasserstein重心为聚合概率测度提供了一种原则性方法，同时保持了其所在空间的几何结构。现有的离散方法由于假设能够获取输入测度的完整样本集，因此缺乏可扩展性。与此同时，神经网络方法虽然具有良好的扩展性，但依赖于复杂的优化问题，且难以有效整合标签信息。我们通过在概率测度空间中的梯度流来解决这些局限性。通过时间离散化，我们实现了一种可扩展算法，该算法具有以下特点：i) 基于小批量最优传输；ii) 通过任务感知函数实现模块化正则化；iii) 将监督信息无缝整合到基础成本函数中。我们在涵盖计算机视觉、神经科学和化学工程的领域自适应基准测试中实证验证了所提出的方法。我们的方法确立了一种新的最先进重心求解器，其中带标签的重心始终优于无标签重心。