Learning conditional distributions on continuous spaces
Cyril Bénézet, Ziteng Cheng, Sebastian Jaimungal
TL;DR
The paper addresses learning conditional distributions $P_x$ in continuous spaces by framing them as measure-valued mappings on unit boxes and proposing two clustering-based estimators, the $\text{r-box}$ and $\text{$k$-NN}$ estimators. It derives convergence and concentration guarantees for these estimators in the Wasserstein metric, with distinct rates depending on the input/output dimensions and the estimator type, and identifies optimal choices for the radius $r$ and neighbor count $k$. To translate these nonparametric estimators into scalable practice, the authors introduce a Lipschitz-adaptive neural network $\tilde P^\Theta$ trained to mimic the empirical estimator, leveraging Approximate NN Search (ANNS-RBSP), the Sinkhorn algorithm for efficient Wasserstein computation, and a convex-potential layer to enable local Lipschitz adaptivity. Empirical results on synthetic 1D and 3D data show LipNet achieving smoother, more accurate conditional distributions and automatic local Lipschitz adaptation, with reproducible code provided. The proposed framework has potential applications in model-based RL and risk-sensitive MDPs where conditional distributions, not just expectations, are essential for decision-making.
Abstract
We investigate sample-based learning of conditional distributions on multi-dimensional unit boxes, allowing for different dimensions of the feature and target spaces. Our approach involves clustering data near varying query points in the feature space to create empirical measures in the target space. We employ two distinct clustering schemes: one based on a fixed-radius ball and the other on nearest neighbors. We establish upper bounds for the convergence rates of both methods and, from these bounds, deduce optimal configurations for the radius and the number of neighbors. We propose to incorporate the nearest neighbors method into neural network training, as our empirical analysis indicates it has better performance in practice. For efficiency, our training process utilizes approximate nearest neighbors search with random binary space partitioning. Additionally, we employ the Sinkhorn algorithm and a sparsity-enforced transport plan. Our empirical findings demonstrate that, with a suitably designed structure, the neural network has the ability to adapt to a suitable level of Lipschitz continuity locally. For reproducibility, our code is available at \url{https://github.com/zcheng-a/LCD_kNN}.