Generalization error bounds using Wasserstein distances

Generalization error of a learning algorithm characterizes the gap between an algorithm's performance on test data versus performance on training data. In recent work, Xu & Raginsky [1] showed that generalization error may be upper- bounded using the mutual information I(S;W) between the input S and the output W of an algorithm. In this paper, we derive upper bounds on the generalization error in terms of a certain Wasserstein distance involving the distributions of S and W under the assumption of a Lipschitz continuous loss function. Unlike mutual information-based bounds, these new bounds are useful even for deterministic learning algorithms, or for algorithms such as stochastic gradient descent. Moreover, we show that in some natural cases these bounds are tighter than mutual information-based bounds.