End-to-End Deep Learning for Predicting Metric Space-Valued Outputs
It addresses a fundamental challenge in machine learning for applications requiring geometry-aware predictions, offering a flexible solution for domains like mortality analysis and network modeling.
The paper tackles the problem of predicting structured, non-Euclidean outputs like probability distributions and networks by introducing E2M, an end-to-end deep learning framework that uses weighted Fréchet means to preserve intrinsic geometry, achieving state-of-the-art performance in simulations and applications.
Many modern applications involve predicting structured, non-Euclidean outputs such as probability distributions, networks, and symmetric positive-definite matrices. These outputs are naturally modeled as elements of general metric spaces, where classical regression techniques that rely on vector space structure no longer apply. We introduce E2M (End-to-End Metric regression), a deep learning framework for predicting metric space-valued outputs. E2M performs prediction via a weighted Fréchet means over training outputs, where the weights are learned by a neural network conditioned on the input. This construction provides a principled mechanism for geometry-aware prediction that avoids surrogate embeddings and restrictive parametric assumptions, while fully preserving the intrinsic geometry of the output space. We establish theoretical guarantees, including a universal approximation theorem that characterizes the expressive capacity of the model and a convergence analysis of the entropy-regularized training objective. Through extensive simulations involving probability distributions, networks, and symmetric positive-definite matrices, we show that E2M consistently achieves state-of-the-art performance, with its advantages becoming more pronounced at larger sample sizes. Applications to human mortality distributions and New York City taxi networks further demonstrate the flexibility and practical utility of the framework.