Improving Optimal Transport using Kernel Methods. (Améliorer le transport optimal à l'aide des méthodes du noyau)

Piyushi Manupriya

Abstract: en fr The Optimal Transport (OT) formulation provides a metric over distributions by lifting the geometry over their support. This thesis focuses on the OT formulation proposed by Kantorovich, which seeks a probabilistic transport plan to optimally couple one distribution with another in the most cost-efficient manner. As comparing distributions is at the core of several machine learning (ML) algorithms, OT has found many applications in ML. While the original optimization is bottlenecked by the computational complexity of solving a linear program, regularizing the objective with the entropy of the transport plan benefits from scalable solvers based on the Sinkhorn algorithm. Such entropy-regularized variants have extended the application of OT to large-scale settings that commonly arise in ML. This thesis begins with a crucial observation that although the geometry induced by OT-based Wasserstein metrics is rich, the sample complexity required to estimate these metrics suffers from the curse of dimensionality. With $m$ as the no. of samples from the measures involved, the minimax estimation error of Wasserstein is known to be $O(m^{-1/d})$, which has an adversarial dependence on the underlying data dimensions. On the other hand, the kernel-based metric, Maximum Mean Discrepancy (MMD), has a dimension-free sample complexity, but the induced geometry is often independent of the kernel employed to capture the same. This leads us to explore if kernel-based regularization can help fix the sample complexity of OT while preserving its metric-related properties. Interestingly, our research results in a new family of metrics with geometry analogous to Wasserstein and a dimension-free sample complexity. We propose novel OT formulations over (potentially unnormalized) measures with kernel-based MMD regularizers to match the transport plan's marginals with the given source and target measures. Furthermore, our careful choice of the squared-MMD metric enables us to exploit the smoothness of the resulting objective to come up with scalable solvers. We derive corresponding theoretical guarantees and showcase improvements to downstream applications. The first challenge we address is that of efficiently estimating OT in high dimensional settings when kernel-based regularization is employed to match the marginals of the transport plan to the given source and target measures. A key contribution of this work is showing that this regularized OT variant not only removes the curse of dimensionality issue in OT but also maintains the metric properties. Our proposed formulation introduces novel statistically efficient metrics belonging to the class of Integral Probability Metrics and obtains an estimation error of $O(m^{-1/2})$. Moreover, we also derive the estimation error when the support of the transport plan is restricted to (a subset of) finite samples available from the source and target measures, as is the case in practical scenarios. We show that, under certain assumptions, the estimation error with this finite parametrization also comes out to be $O(m^{-1/2})$. To the best of our knowledge, such dimension-free error bounds are not known for other variants of regularized OT. Further, leveraging the smoothness of the optimization objective induced by squared-MMD regularization, we demonstrate an accelerated projected gradient descent solver that matches the computational efficiency of the Sinkhorn solver popularly used for entropy-regularized OT. Next, we use the proposed OT formulation with kernel-based regularization to alleviate the issue of non-sparsity in transport plans that result from entropy regularization in OT. While the entropy-regularized OT variants are computationally attractive, the resulting transport plans are dense, with each entry being strictly positive. This, in turn, affects the interpretability of the alignments obtained through the transport plan. We discuss how kernel-based regularization allows us to induce structured sparsity in the finitely parameterized transport plan. Leveraging the smoothness and strong convexity of our optimization objective, we recast the problem of finding a sparse solution to the optimal transport problem to that of a submodular maximization problem. We propose algorithmic solvers, derive approximation guarantees, and show that the duality gap for the problem obtained with the proposed approach is better than that of the existing approaches. We highlight that other popular variants of OT do not enjoy a smooth objective function and, hence, can not benefit from our proposed approach to this problem. This again shows the efficacy of the proposed kernel-based regularization in the OT formulation. Finally, we demonstrate this thesis idea to solve OT problems involving conditional measures, which are implicitly specified through joint samples. Given samples from two joint distributions, we consider the problem of OT between them when conditioned on a common variable. This setting is particularly difficult when the conditioned variable is continuous. The key challenge in estimating OT in this setting comes in enforcing the marginal-matching constraints involving the conditionals, as the samples provided are not from the conditionals but from the joints. Our formulation employs kernelized least-squares-based MMD-regularizers computed over the joint samples to implicitly match the transport plan’s marginals with the empirical conditionals. We also discuss learning the transport plan with parametric models. As the MMD metric is meaningful even with measures having non-overlapping support, this facilitates us in employing implicit generative models to learn the transport plan. Under mild conditions, the error in estimating the divergence between the conditionals decays at $O(m^{-1/4})$ where $m$ is the no. of samples. To the best of our knowledge, ours is the first consistent OT formulation between empirical conditional measures. We empirically validate the theoretical guarantees of our proposed formulations. We show improvements in diverse applications such as domain adaptation, two-sample hypothesis tests, prompt learning for few-shot classification, cell-population dynamics, aligning words in sentence pairs and learning mixture-of-experts models. This thesis underscores the critical role of kernel-based regularization in various formulations of optimal transport. In particular, we demonstrate how the choice of such regularization results in dimension-free estimation error while maintaining the metric properties, helps efficiently induce sparsity structure on the transport plan and results in a consistent formulation for comparing conditional measures. We believe this research will open avenues for future studies on the interplay between OT and kernel methods. La formulation du transport optimal (OT) fournit une métrique sur les distributions en relevant la géométrie au-dessus de leur support. Cette thèse se concentre sur la formulation OT proposée par Kantorovich, qui recherche un plan de transport probabiliste pour coupler de manière optimale une distribution à une autre de la manière la plus rentable. La comparaison des distributions étant au cœur de plusieurs algorithmes d'apprentissage automatique (ML), l'OT a trouvé de nombreuses applications en ML. Alors que l'optimisation initiale est freinée par la complexité de calcul de la résolution d'un programme linéaire, la régularisation de l'objectif avec l'entropie du plan de transport bénéficie de solveurs évolutifs basés sur l'algorithme de Sinkhorn. Ces variantes régularisées par l'entropie ont étendu l'application de l'OT aux contextes à grande échelle fréquemment rencontrés en ML. Cette thèse part d'une observation cruciale : bien que la géométrie induite par les métriques de Wasserstein basées sur l'OT soit riche, la complexité de l'échantillon nécessaire à son estimation souffre du fléau de la dimensionnalité. Avec $m$ comme nombre d'échantillons des mesures concernées, l'erreur d'estimation minimax de Wasserstein est connue pour être de $O(m^{-1/d})$, ce qui présente une dépendance antagoniste aux dimensions des données sous-jacentes. En revanche, la métrique basée sur le noyau, l'écart moyen maximal (MMD), présente une complexité d'échantillon sans dimension, mais la géométrie induite est souvent indépendante du noyau utilisé pour la capturer. Cela nous amène à explorer si la régularisation basée sur le noyau peut contribuer à corriger la complexité de l'échantillon d'OT tout en préservant ses propriétés métriques. Il est intéressant de noter que nos recherches aboutissent à une nouvelle famille de métriques avec une géométrie analogue à celle de Wasserstein et une complexité d'échantillon sans dimension. Nous proposons de nouvelles formulations OT sur des mesures (potentiellement non normalisées) avec des régularisateurs MMD à noyau afin de faire correspondre les marges du plan de transport aux mesures source et cible données. De plus, notre choix judicieux de la métrique MMD au carré nous permet d'exploiter la régularité de l'objectif résultant pour obtenir des solveurs évolutifs. Nous en déduisons les garanties théoriques correspondantes et présentons des améliorations pour les applications en aval. Le premier défi que nous abordons est l'estimation efficace de l'OT dans des contextes de grande dimension, lorsque la régularisation par noyau est utilisée pour faire correspondre les marges du plan de transport aux mesures source et cible données. Une contribution clé de ce travail est de montrer que cette variante régularisée de l'OT non seulement supprime le problème de dimensionnalité en OT, mais préserve également les propriétés métriques. Notre formulation proposée introduit de nouvelles métriques statistiquement efficaces appartenant à la classe des métriques de probabilité intégrale et obtient une erreur d'estimation de $O(m^{-1/2})$. De plus, nous dérivons également l'erreur d'estimation lorsque le support du plan de transport est restreint à (un sous-ensemble d') échantillons finis disponibles à partir des mesures source et cible, comme c'est le cas dans des scénarios pratiques. Nous montrons que, sous certaines hypothèses, l'erreur d'estimation avec cette paramétrisation finie est également de $O(m^{-1/2})$. À notre connaissance, de telles bornes d'erreur sans dimension ne sont pas connues pour d'autres variantes de l'OT régularisé. De plus, en tirant parti de la régularité de l'objectif d'optimisation induit par la régularisation MMD au carré, nous démontrons un solveur de descente de gradient projeté accéléré qui égale l'efficacité de calcul du solveur Sinkhorn couramment utilisé pour l'OT régularisé par entropie. Nous utilisons ensuite la formulation OT proposée avec une régularisation par noyau pour pallier le problème de parcimonie dans les plans de transport résultant de la régularisation entropique en OT. Si les variantes OT régularisées par entropie sont intéressantes sur le plan informatique, les plans de transport résultants sont denses, chaque entrée étant strictement positive. Ceci, à son tour, affecte l'interprétabilité des alignements obtenus grâce au plan de transport. Nous expliquons comment la régularisation par noyau permet d'induire une parcimonie structurée dans le plan de transport à paramètres finis. En tirant parti de la régularité et de la forte convexité de notre objectif d'optimisation, nous transformons le problème de recherche d'une solution parcimonieuse au problème de transport optimal en un problème de maximisation sous-modulaire. Nous proposons des solveurs algorithmiques, déduisons des garanties d'approximation et montrons que l'écart de dualité pour le problème obtenu avec l'approche proposée est meilleur que celui des approches existantes. Nous soulignons que d'autres variantes populaires de l'OT ne bénéficient pas d'une fonction objectif lisse et, par conséquent, ne peuvent pas bénéficier de notre approche proposée pour ce problème. Cela démontre une fois de plus l'efficacité de la régularisation par noyau proposée dans la formulation de l'OT. Enfin, nous démontrons cette idée de thèse pour résoudre des problèmes d'OT impliquant des mesures conditionnelles, implicitement spécifiées par des échantillons joints. Étant donnés des échantillons de deux distributions jointes, nous étudions le problème de l'OT entre elles conditionnées par une variable commune. Ce contexte est particulièrement complexe lorsque la variable conditionnée est continue. Le principal défi pour estimer l'OT dans ce contexte réside dans l'application des contraintes de correspondance marginale impliquant les conditions, car les échantillons fournis ne proviennent pas des conditions, mais des articulations. Notre formulation utilise des régularisateurs MMD à noyaux basés sur les moindres carrés, calculés sur les échantillons joints, pour faire correspondre implicitement les marginales du plan de transport aux conditions empiriques. Nous abordons également l'apprentissage du plan de transport avec des modèles paramétriques. La métrique MMD étant pertinente même avec des mesures dont le support ne se chevauche pas, cela facilite l'utilisation de modèles génératifs implicites pour l'apprentissage du plan de transport. Dans des conditions modérées, l'erreur d'estimation de la divergence entre les mesures conditionnelles diminue à $O(m^{-1/4})$, où $m$ est le nombre d'échantillons. À notre connaissance, notre formulation OT est la première cohérente entre des mesures conditionnelles empiriques. Nous validons empiriquement les garanties théoriques de nos formulations proposées. Nous montrons des améliorations dans diverses applications telles que l'adaptation de domaine, les tests d'hypothèses à deux échantillons, l'apprentissage rapide pour la classification à quelques coups, la dynamique des populations cellulaires, l'alignement de mots dans des paires de phrases et l'apprentissage de modèles mixtes d'experts. Cette thèse souligne le rôle crucial de la régularisation par noyau dans diverses formulations du transport optimal. Nous démontrons notamment comment le choix d'une telle régularisation génère une erreur d'estimation adimensionnelle tout en préservant les propriétés métriques, contribue à induire efficacement une structure parcimonieuse sur le plan de transport et conduit à une formulation cohérente pour la comparaison des mesures conditionnelles. Nous pensons que cette recherche ouvrira des perspectives pour de futures études sur l'interaction entre les méthodes OT et par noyau.

External IDs:dblp:phd/hal/Manupriya25