Wasserstein Distributional Learning

TL;DR

Wasserstein Distributional Learning (WDL) introduces a novel framework for modeling conditional densities using Wasserstein distance, overcoming traditional constraints and capturing nonlinear dependencies more effectively.

Contribution

The paper proposes WDL, a flexible density-on-scalar regression framework with a new model class and an efficient optimization algorithm, advancing the modeling of conditional densities.

Findings

WDL outperforms existing methods in characterizing nonlinear density dependencies.

The framework effectively captures complex distributional relationships in simulations.

Real-world applications demonstrate WDL's practical utility and improved accuracy.

Abstract

Learning conditional densities and identifying factors that influence the entire distribution are vital tasks in data-driven applications. Conventional approaches work mostly with summary statistics, and are hence inadequate for a comprehensive investigation. Recently, there have been developments on functional regression methods to model density curves as functional outcomes. A major challenge for developing such models lies in the inherent constraint of non-negativity and unit integral for the functional space of density outcomes. To overcome this fundamental issue, we propose Wasserstein Distributional Learning (WDL), a flexible density-on-scalar regression modeling framework that starts with the Wasserstein distance $W_2$ as a proper metric for the space of density outcomes. We then introduce a heterogeneous and flexible class of Semi-parametric Conditional Gaussian Mixture Models (SCGMM) as the model class $\mathfrak{F} \otimes \mathcal{T}$. The resulting metric space $(\mathfrak{F} \otimes \mathcal{T}, W_2)$ satisfies the required constraints and offers a dense and closed functional subspace. For fitting the proposed model, we further develop an efficient algorithm based on Majorization-Minimization optimization with boosted trees. Compared with methods in the previous literature, WDL better characterizes and uncovers the nonlinear dependence of the conditional densities, and their derived summary statistics. We demonstrate the effectiveness of the WDL framework through simulations and real-world applications.