Representation Learning for Equivariant Inference with Guarantees

In many real-world applications of regression, conditional probability estimation, and uncertainty quantification, exploiting symmetries rooted in physics or geometry can dramatically improve generalization and sample efficiency. While geometric deep learning has made empirical advances by incorporating symmetry and geometry priors, less attention has been given to statistical learning guarantees. In this paper, we introduce an equivariant representation learning framework that simultaneously addresses regression, conditional probability estimation, and uncertainty quantification while providing first-of-its-kind non-asymptotic statistical learning guarantees. Grounded in operator and group representation theory, our framework approximates the spectral decomposition of the conditional expectation operator, building representations that are both equivariant and disentangled along independent symmetry quotient groups. Empirical evaluations on synthetic datasets and real-world robotics applications confirm the potential of our approach, matching or outperforming existing equivariant baselines in regression while providing well-calibrated uncertainty estimates.

Many machine learning problems require more than a single prediction: they also require estimating probabilities and uncertainty. In many real-world settings, the data has known symmetries coming from physics or geometry. For example, a system may behave consistently under rotations, reflections, or other transformations. Using these symmetries can help models generalize better and learn from fewer examples. Geometric deep learning has made strong empirical progress by building neural networks that respect such symmetries, but there is still limited understanding of how symmetry-aware representations can be learned with statistical guarantees. We introduce eNCP, a representation learning framework that uses known symmetries to support regression, conditional probability estimation, and uncertainty quantification within a single method. The framework is based on a mathematical view of conditional prediction as an operator, and it learns representations that are both symmetry-respecting and disentangled according to independent symmetry components. We prove non-asymptotic learning guarantees linking representation quality to sample complexity. Experiments on synthetic and real-world datasets show that eNCP matches or improves over existing equivariant methods while producing well-calibrated uncertainty estimates.