[Paper Review] Wasserstein Distributional Robustness and Regularization in Statistical Learning.

Rui Gao, Xi Chen+1 more|arXiv (Cornell University)|Dec 17, 2017

Risk and Portfolio Optimization48 references89 citations

TL;DR

This paper proposes a distributionally robust optimization framework using Wasserstein distance to enhance generalization in statistical learning. It establishes that Wasserstein distributional robustness is asymptotically equivalent to regularization with a gradient-norm penalty, offering a principled approach to regularizing high-dimensional, non-convex problems, including in deep learning via Wasserstein GANs.

ABSTRACT

A central question in statistical learning is to design algorithms that not only perform well on training data, but also generalize to new and unseen data. In this paper, we tackle this question by formulating a distributionally robust stochastic optimization (DRSO) problem, which seeks a solution that minimizes the worst-case expected loss over a family of distributions that are close to the empirical distribution in Wasserstein distances. We establish a connection between such Wasserstein DRSO and regularization. More precisely, we identify a broad class of loss functions, for which the Wasserstein DRSO is asymptotically equivalent to a regularization problem with a gradient-norm penalty. Such relation provides new interpretations for problems involving regularization, including a great number of statistical learning problems and discrete choice models (e.g. multinomial logit). The connection suggests a principled way to regularize high-dimensional, non-convex problems. This is demonstrated through the training of Wasserstein generative adversarial networks in deep learning.

Motivation & Objective

• To address the challenge of generalization in statistical learning beyond training data.
• To develop a principled framework for robust optimization under distributional uncertainty.
• To establish a theoretical connection between distributionally robust optimization and regularization.
• To provide new interpretations for regularization in models like multinomial logit and deep neural networks.
• To demonstrate the practical utility of the framework in deep learning, particularly in training Wasserstein GANs.

Proposed method

• Formulates a distributionally robust stochastic optimization (DRSO) problem minimizing worst-case expected loss over distributions within a Wasserstein ball around the empirical distribution.
• Identifies a broad class of loss functions for which the DRSO problem is asymptotically equivalent to a regularization problem with a gradient-norm penalty.
• Uses tools from optimal transport and empirical process theory to derive the asymptotic equivalence between DRSO and regularization.
• Applies the theoretical framework to deep learning by showing its relevance in training Wasserstein GANs.
• Demonstrates that the robustness induced by the Wasserstein DRSO naturally leads to implicit regularization in high-dimensional, non-convex settings.
• Provides a unifying perspective that interprets existing regularization techniques as arising from distributional robustness under Wasserstein metrics.

Experimental results

Research questions

• RQ1 How can distributional robustness via Wasserstein distance improve generalization in statistical learning?
• RQ2 What is the theoretical connection between distributionally robust optimization and regularization?
• RQ3 For which classes of loss functions does distributional robustness under Wasserstein distance lead to gradient-norm regularization?
• RQ4 Can this framework be applied to non-convex, high-dimensional problems such as deep neural networks?
• RQ5 How does this approach enhance training stability and performance in generative modeling, such as in Wasserstein GANs?

Key findings

• The Wasserstein DRSO problem is asymptotically equivalent to a regularization problem with a gradient-norm penalty for a broad class of loss functions.
• This equivalence provides a principled interpretation of regularization as a form of distributional robustness under the Wasserstein metric.
• The framework offers a new theoretical foundation for understanding and designing regularization in discrete choice models, such as multinomial logit.
• The approach enables robust generalization in high-dimensional, non-convex problems by implicitly regularizing the model's gradient behavior.
• The method improves training stability and performance in deep learning, as demonstrated in the training of Wasserstein GANs.
• The theoretical results suggest that distributional robustness via Wasserstein distance naturally induces regularization, enhancing model generalization.