Representation Learning for Equivariant Inference with Guarantees

Code repository for ICML 2026 submission "Representation Learning for Equivariant Inference with Guarantees".

Installation

The repository holds a python package named symm_rep_learn that implements the models and baselines described in the paper, with ready-to-import modules for third party development. To install the core library do:

pip install -e .

Additional depdendencies for plotting, logging and visualization are needed to run the experiments and notebooks in the paper/ folder. If you want to run these, please install the optional dependencies with:

pip install -e ".[paper]"

Repository Structure

.
├── symm_rep_learn/              # Core library
│   ├── inference/               # Inference-time modules (e.g., cCDF estimation, conditional quantile regression)
│   ├── models/                  # Model implementations  (eNCP, NCP, baselines)
│   ├── nn/                      # Neural network components
│   └── mysc/                    # Utilities and theory
├── paper/                       # Paper experiments and examples
    ├── examples/                # Reproducible examples
    ├── experiments/             # Main experiments
    ├── plots/                   # Generated plots and figures
    └── results/                 # Experimental results

Ready-to-use models

• ENCP (Equivariant Neural Conditional Probability) extends the neural conditional operator with equivariant embeddings and statistics so that conditional expectations respect symmetry constraints.
• NCP (Neural Conditional Probability) is the base operator that learns low-rank factorizations of conditional expectations and exposes helpers such as conditional_expectation and fit_linear_decoder for regression tasks.

Inference modules

• ENCPConditionalCDF Uses a trained ENCP to perform symmetry-aware conditional CDF (cCDF) estimation and conditional quantile regression.
• NCPConditionalCDF Uses a trained NCP to perform cCDF estimation and conditional quantile regression.

Reproducible examples

1. Conditional expectation (regression) with uncertainty quantification

We demonstrate conditional expectation (regression) with uncertainty quantification in the notebook conditional_expectation_regression_1D.ipynb. The notebook tackles a picewise 1D regression where we aim to predict both the expected value of Y given X and confidence intervals (lower and upper quantiles) for the prediction. Confidence intervals are of paramount importance in the regions where the conditional distribution $\mathbb{P}(y \mid x)$ is multimodal or skewed.

The notebook illustrates how to use the eNCP and NCP models to estimate conditional expectations (regression) and conditional quantiles (uncertainty quantification), without any retraining needed for estimation of quantiles of different coverage levels.

2. Conditional quantile regression

The notebook conditional_quantile_regression_synthetic.ipynb shows how the eNCP and NCP framework model conditional probabilities enabling the prediction of the conditional Cumulative Distribution Function (cCDF) enabling the regression of conditional quantiles of any desired coverage level.

The results show how the eNCP and NCP models outperform training frameworks that aim to directly predict quantiles of a fixed coverage level (i.e., Conditional Quantile Regression, CQR)

3. Uncertainty quantification in ground reaction force estimation in legged locmotion

The notebook conditional_quantile_regression_quadruped.ipynb shows how the eNCP and NCP framework can be used for uncertainty quantification in the estimation of ground reaction forces (GRFs) in quadruped locomotion over rough terrain.

4. Sensitivity analysis to symmetry misspecification

The notebook misspecified_sensitivity_analysis_1D.ipynb investigates robustness to symmetry-prior misspecification in the same 1D synthetic setting used in experiment (1). In particular, we study the cases of incorrect and extrinsic misspecification, following the taxonomy of Wang et al. (2023).

1. Extrinsic misspecification of $P_{\mathbf{x}}$ $\mathbb{G}$-invariance:
   This scenario arises when the assumption that the marginal distribution of $\mathbf{x}$ is $\mathbb{G}$-invariant is violated. Specifically, the support of the random variable in the training, validation, and test sets is not $\mathbb{G}$-invariant, so the empirical distribution of $\mathbf{x}$ is biased toward only a subset of the support of $P_{\mathbf{x}}$. We study this setting by training on a biased half-space ($\mathbf{x} > 0$) and evaluating performance both on the same biased support and on the full support ($\mathbf{x} > 0$ and $\mathbf{x} \in \mathbb{R}$).

   

   In this setting, assuming $\mathbb{G}$-invariance of $P_{\mathbf{x}}$ and using the eNCP model acts as a form of regularization, enabling out-of-distribution generalization without degrading performance on the biased support.

2. Incorrect misspecification of $P_{\mathbf{y} \mid \mathbf{x}}$ $\mathbb{G}$-invariance:
   This scenario corresponds to violations of the assumption that the conditional distribution of $\mathbf{y}$ given $\mathbf{x}$ is $\mathbb{G}$-equivariant on a subset of the support of $\mathbf{x}$. We consider two types of incorrect misspecification:

   • Unbiased incorrect misspecification: Here, we progressively scale the heteroscedastic noise amplitude in the region $\mathbf{x} > 1$. This violates the $\mathbb{G}$-invariance of $P_{\mathbf{y} \mid \mathbf{x}}$ while preserving the $\mathbb{G}$-equivariance of the conditional expectation $\mathbb{E}[\mathbf{y} \mid \mathbf{x}]$.
   
   • Biased incorrect misspecification: Here, we progressively introduce a linear bias in the region $\mathbf{x} > 1$. This violates both the $\mathbb{G}$-invariance of $P_{\mathbf{y} \mid \mathbf{x}}$ and the $\mathbb{G}$-equivariance of the conditional expectation $\mathbb{E}[\mathbf{y} \mid \mathbf{x}]$ on the subset $|\mathbf{x}| > 1$.
   

The results, show that the performance of the eNCP model deteriorates continuously with the degree of misspecification

Baseline implementations

• CQR (Conditional Quantile Regression) implements the standard two-head pinball-loss baseline.
• eCQR (Equivariant Conditional Quantile Regression) wraps equivariant MLPs to enforce symmetry-aware prediction intervals.
• DRF (Density Ratio Fitting) provides density-ratio based estimators of pointwise mutual dependency.
• InvDRF (Invariant Density Ratio Fitting) adapts DRF to invariant equivariant modules for symmetry-preserving density-ratio estimation.