Incorporating Importance Weighting in Optimal Transport Based Domain Alignment

Domain adaptation theory studies upper bounds on the target risk in order to mitigate performance loss of machine learning models due to distribution shift. In this paper, we take a closer look at the optimization of one such bound based on optimal transport (OT) and propose various strategies that improve the optimization in practice. We first introduce *gradual shift* and *probabilistic margin* assumptions to control the incomputable entanglement term that appears in the bounds. We prove that under these assumptions, better optimization of the computable part of the bound can translate to better target accuracies. Motivated by this fact, we tighten the bound, via importance weighting of the source (output) distribution, to obtain the *weighted* Wasserstein regularized risk ($\mathrm{W}^2\mathrm{R}^{2}$), that is often easier to minimize than the original bound. $\mathrm{W}^2\mathrm{R}^{2}$ is shown to be equivalent to an unbalanced OT problem, which in the limit converges to a nearest neighbor based alignment strategy. We highlight the tradeoffs faced with such an approach and show that a suitably regularized $\mathrm{W}^2\mathrm{R}^{2}$ improves over the state of the art and is robust to multiple distribution shifts under different models, confirming, moreover, the validity of our assumptions.

Whenever machine learning models are deployed in new environments, sufficiently distinct from the ones they are trained on, their performance deteriorates rapidly. Domain adaptation studies ways to maintain or improve the accuracy of the models in these new environments, also called domains. One standard way to do so is via feature matching - which tries to align the features from different environments over the layers of a neural network, but this procedure can unnecessarily constrain the models and even at times fail to improve the performance. Instead one can perform a relaxed alignment, that does not try to match the features exactly but only makes sure the target features are within the natural variability of the original environment's features. Building on this idea, we propose a more relaxed alignment strategy in this paper and show that it is closely connected to a technique called unbalanced optimal transport. We also explore assumptions under which such alignment is guaranteed to produce models with high accuracy on the new environment. Finally, we verify in experiments that this more relaxed alignment improves over the standard one.