Equivariant Representation Learning for Symmetry-Aware Inference with Guarantees

TL;DR

This paper introduces an equivariant representation learning framework that leverages symmetry in data to improve regression, probability estimation, and uncertainty quantification, with theoretical guarantees and strong empirical results.

Contribution

It presents the first framework combining equivariant representations with non-asymptotic statistical guarantees for regression and uncertainty quantification.

Findings

Outperforms existing equivariant methods in regression tasks

Provides well-calibrated parametric uncertainty estimates

Works effectively on synthetic and real-world robotics datasets

Abstract

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 significant empirical advances by incorporating group-theoretic structure, 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 subgroups. 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 additionally providing well-calibrated parametric uncertainty estimates.