Functional Diversity in Deep Ensembles via Wasserstein Representation Alignment

Deep Ensembles (DE) are widely recognized for their robustness and improved uncertainty quantification in machine learning tasks. However, a persistent challenge lies in ensuring sufficient diversity among ensemble members: independently trained models often converge to similar solutions, limiting the ensemble’s overall effectiveness. In this work, we propose a novel Wasserstein regularization approach in function space. Specifically, we maximize pairwise Wasserstein or heuristic Gromov–Wasserstein (GW) distances between hidden layers’ representations of ensemble members. By interpreting these representations as empirical probability distributions, our method leverages the geometry of optimal transport to enforce functional diversity. Furthermore, we demonstrate that regularizing in function space—via distributional alignment using Wasserstein-based metrics—yields better diversity and generalization than approaches based on weight space alone. Empirical results on benchmark datasets demonstrate that our method outperforms standard DE and weight-space regularized variants in both predictive accuracy and uncertainty estimation. Code and experiments are available at https://github.com/leotaku98/wass_de .