Equivariant Covariance Tensors: Guaranteed SPD Uncertainty for Tensor-Valued Geometric Learning

Ruihan Liu; Yu Ji; Jianbo Yu; Shifu Yan; Qingchao Jiang

cs.LGAug 25, 2026

Equivariant Covariance Tensors: Guaranteed SPD Uncertainty for Tensor-Valued Geometric Learning

Authors: Ruihan Liu, Yu Ji, Jianbo Yu, Shifu Yan, Qingchao Jiang

Organizations: College of Intelligent Robotics and Advanced Manufacturing, Fudan University, Shanghai 200433, China · School of Microelectronics, Fudan University, Shanghai 200433, China · ByteDance, Beijing, China · School of Information Science and Engineering, East China University of Science and Technology, Shanghai 200237, China

Abstract

Tensor-valued prediction is fundamental to geometric deep learning, yet uncertainty quantification (UQ) for such outputs remains an open challenge. While E(3)-equivariant neural networks excel at point estimates, they lack rigorous confidence measures. We focus on symmetric rank-2 tensor prediction, where the target has six Kelvin--Mandel coordinates and full uncertainty is represented by a $6\times6$ covariance matrix. We introduce a framework for E(3)-equivariant UQ, modeling the full predictive distribution where both mean and covariance preserve rotational symmetry. Our approach decomposes the covariance into irreducible representations $\mathrm{Sym}^2(ρ_c) \cong 2\times(l=0) \oplus 2\times(l=2) \oplus 1\times(l=4)$ . By mapping from the flat Lie algebra $\mathfrak{sym}(6)$ to the curved SPD manifold via matrix exponentiation, we strictly ensure positive-definite covariances while maintaining exact equivariance. Furthermore, we formulate a Log-Euclidean Equivariant Scoring Objective (LE-ESO)---a robust surrogate loss based on the Multivariate Laplace distribution---providing robustness to heavy-tailed errors and stable optimization. Validation on ModelNet40 inertia tensors and Materials Project dielectric tensors demonstrates that our method achieves competitive performance and provides physically consistent, symmetry-preserving uncertainty estimates with useful risk and OOD sensitivity.

Figures & tables

Figure 1: E(3)-equivariant tensor uncertainty framework. (1) Feature Decomposition: covariance head models symmetry via irreducible representations ( $\ell=0,2,4$ ). (2) Manifold Mapping: network predicts unconstrained $A\in\mathfrak{sym}(6)$ . (3) Validity Constraint: $\Sigma=\exp(A)$ projects to SPD manifold $\mathcal{P}_{6}$ .

Figure 2: Geometric interpretation of equivariant covariance construction. The network operates in the Lie algebra $\mathfrak{g}\cong\mathbb{R}_{\mathrm{sym}}^{6\times 6}$ (left), producing $A(X)$ . The exponential map $\Sigma=\exp(A)$ (center) projects onto the SPD manifold $\mathcal{P}_{6}$ (right), guaranteeing valid uncertainties while preserving $E(3)$ -equivariance.

Figure 3: 3D uncertainty visualization on ModelNet40. (Top) Point clouds with uncertainty ellipsoids (blue) and principal axes: predicted (red) vs ground truth (green). (Bottom) Anisotropic covariance matrices. Samples show varying errors (0.293–0.472) and Mahalanobis distances (0.37–0.50), demonstrating geometric-adaptive uncertainty that rotates with the object shape.

Method	MAE	RMSE	LE-ESO	SPD
Ours (Full Cov)	$\mathbf{0.078}$	$\mathbf{0.128}$	$\mathbf{10.66}$	✓
Diagonal Cov	0.092	0.156	12.84	✓
Deterministic (MSE)	0.083	0.135	N/A	✗

Table 1: ModelNet40 inertia tensor prediction. Our equivariant full-covariance framework achieves strong performance with exact geometric consistency and physical constraints.

Figure 4: Performance on Materials Project. (a) Parity plot of diagonal components ( $R^{2}=0.659$ ). (b) Multivariate Laplace reliability diagram (MACE $=0.0489$ ).

Method	Type	MAE ( $\downarrow$ )	LE-ESO ( $\downarrow$ )	ES ( $\downarrow$ )	SPD?
SEConv ( Heilman et al., 2024 )	Point	4.702	N/A	N/A	N/A
DTNet ( Mao et al., 2024 )	Point	1.91	N/A	N/A	N/A
GoeCTP ( Hua et al., 2026 )	Point	1.41	N/A	N/A	N/A
Deterministic (MACE)	Point	2.10	N/A	N/A	N/A
MACE Ensemble ( $N=5$ )	UQ	1.96	-2.55	0.78	$\approx$
Diagonal UQ	UQ	2.25	-1.85	0.87	✓

Table 2: Performance on Materials Project dielectric dataset. GoeCTP ( Hua et al., 2026 ) employs a scalable equivariant architecture. DTNet ( Mao et al., 2024 ) uses universal potential embeddings. MACE-Ens. is a 5-model deep ensemble. Ours (Calib.) applies temperature scaling (

T\approx 0.05

). ES generalizes CRPS to multivariate settings.

Figure 5: Uncertainty diagnostics and utility analysis. (a) Sharpness distribution of 95% confidence volumes, demonstrating the model’s ability to assign heteroscedastic uncertainties. (b) Risk-coverage analysis comparing ranking by directional uncertainty (

\lambda_{\max}(\Sigma)

) and total uncertainty (

\mathrm{Tr}(\Sigma)

). Directional uncertainty provides a modest but consistent advantage in identifying high-error samples.

Training Objective	MAE ( $\downarrow$ )	MACE ( $\downarrow$ )	ES ( $\downarrow$ )
Gaussian NLL	1.78	0.092	1.05
Energy Score (ES)	1.64	0.054	0.81
LE-ESO (Ours)	1.55	0.049	0.66

Table 3: Ablation of training objectives on Materials Project. All variants use the same E(3)-equivariant backbone and matrix-exponential head. LE-ESO (Ours) demonstrates superior robustness to heavy-tailed noise compared to Gaussian NLL.

Method	$\mathcal{E}_{\mu}$	$\mathcal{E}_{\Sigma}$	Mean equiv.?	Cov. equiv.?	SPD?
Baseline A (non-equivariant GNN + Cholesky)	$1.36\times 10^{0}$	$1.07\times 10^{0}$	✗	✗	✓
Baseline B (E3NN + coordinate-wise heads)	$1.28\times 10^{0}$	$1.03\times 10^{0}$	✗	✗	✓
Baseline B ′ (equivariant mean + Cholesky cov.)	$1.43\times 10^{-6}$	$4.30\times 10^{-1}$	✓	✗	✓
Baseline C (E3NN + direct equivariant cov.)	$7.59\times 10^{-7}$	$4.76\times 10^{-7}$	✓	✓	✗
Ours (E3NN + matrix exp.)	$\mathbf{2.39\times 10^{-7}}$	$\mathbf{2.75\times 10^{-7}}$	✓	✓	✓

Table 4: Equivariance and SPD-validity analysis. Baseline B ′ isolates the Cholesky covariance failure mode (equivariant

\mu

, non-equivariant

\Sigma

). Only the matrix-exponential head achieves both covariance equivariance and strict SPD validity.

Figure 6: Chemical OOD sensitivity analysis. The predicted uncertainty increases as the crystal lattice is populated with unseen chemical species, suggesting that the model provides a useful distribution-shift risk signal.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Property	Gaussian NLL	Laplace NLL (LE-ESO)
Mahalanobis term	$\frac{1}{2}D_{M}^{2}$	$D_{M}$
Penalty shape	Quadratic	Linear
Log-det coefficient	$\frac{1}{2}$	$1$ (tunable as $\alpha$ )
Outlier sensitivity	High (quadratic growth)	Low (linear growth)
Gradient for large $D_{M}$	$\propto D_{M}$	$\propto 1$
Statistical assumption	Light-tailed errors	Heavy-tailed errors

Appendix

Table 5: Comparison of Gaussian and Laplace negative log-likelihood formulations.

Method	MAE $\downarrow$	RMSE $\downarrow$	Mean Tensor PSD Rate
Deterministic (MSE)	0.0333	0.0616	99.2%
Full UQ (Ours)	0.0346	0.0633	99.5%

Appendix

Table 6: ModelNet40 shape-covariance validation. The goal is controlled geometric validation rather than replacing the closed-form estimator. “Mean Tensor PSD Rate” refers to the fraction of predicted mean shape-covariance tensors that are positive semi-definite, distinct from the SPD validity of the predictive covariance

\Sigma(X)

Model	Time / Batch	Relative Cost
Deterministic	130 ms	1.00 $\times$
Diagonal UQ	132 ms	1.015 $\times$
Full Covariance UQ	570 ms	4.4 $\times$

Appendix

Table 7: Runtime profiling on Materials Project (RTX 4060 Ti, batch size 16). The full-covariance model is more expensive primarily because of the covariance branch and its backpropagation, rather than the matrix exponential alone.

$\alpha$	Best val MAE $\downarrow$	Best val LE-ESO $\downarrow$
0.03	0.3634	0.7911
0.10	0.4176	0.7165
0.30	0.4566	0.3469
1.00	0.3519	-2.9393

Appendix

Table 8: Sensitivity to the LE-ESO weight

\alpha

on the Materials Project dielectric task. MAE is measured in log-Kelvin–Mandel space and is comparable across rows. The LE-ESO value is evaluated with the same

\alpha

used for training, so it reflects the optimized objective for each setting rather than a fixed cross-

\alpha

NLL.

Figure 7: Training dynamics. Two-stage optimization: warmup (5 epochs) then LE-ESO. (a) LE-ESO convergence. (b) MAE stability for diagonal/off-diagonal components. (c) Balance between data fit (

\mathbb{E}[D_{M}]

) and regularization (

\log\det\Sigma

Panel of Figure 8 — Figure 8: Numerical stability analysis. (a) Positive eigenvalues ensure SPD validity. (b) Moderate condition numbers indicate numerical stability.

Figure 9: SPD validity on inertia-tensor task. (a) Minimum eigenvalue distribution (all positive). (b) $\log_{10}$ condition numbers (well-conditioned). (c) Total uncertainty $\mathrm{Tr}(\Sigma)$ .

Uncertainty Quantification

Uncertainty quantification estimates how reliable predictions and scientific inferences are, and how uncertainty propagates into decisions.

See all 660 papers

May 26, 2025

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.

Daniel Ordoñez-Apraez, Vladimir Kostić, Alek Fröhlich +3

CSML, Italian Institute of Technology, Genova, Italy · Universit`a di Genova, Genova, Italy · University of Novi Sad, Novi Sad, Serbia +2

Disentangled Representation Learning Representation Learning Uncertainty Quantification Conditional Distribution Estimation Equivariant Representation Learning Statistical Learning Theory

Ask PDF

Jun 4, 2026

Equivariant Neural Belief Propagation

Probabilistic inference over spatially embedded variables requires beliefs that respect

SE(3)

symmetry, yet existing equivariant networks produce only scalars and vectors -- not the rank-2 precision tensors needed for anisotropic uncertainty, and single-component messages collapse multi-modal energy landscapes to physically meaningless averages. We introduce Equivariant Neural Belief Propagation (ENBP), a factor-graph framework whose messages are equivariant Gaussian mixture models with sufficient statistics that transform exactly under

SE(3)

. Rank-2 precision matrices are synthesised via equivariant outer products, ingested through differentiable spectral decomposition, and kept tractable by a greedy KL-based mixture reduction that provably commutes with

SE(3)

. On GEOM-QM9 and GEOM-Drugs, ENBP achieves 98.9% conformational coverage at 0.090

\mathring{A}

error with sub-second latency -- over

100\times

faster than diffusion baselines at higher accuracy. On multi-body robotic inference, vanilla loopy BP diverges at 15+ agents while ENBP converges with near-zero collision rates and machine-precision equivariance error (

{\sim}10^{-7}

vs.\

10^{-1}

for augmented baselines).

Zehua Cheng, Wei Dai, Jiahao Sun

Department of Computer Science, University of Oxford, Oxford, United Kingdom · FLock.io, London, United Kingdom

Equivariant Representation Learning Factor Graph Optimization Gaussian Mixture Models Bayesian Inference Equivariant Neural Networks Probabilistic Graphical Models+2

Ask PDF

May 19, 2026

Exact Symmetry as Algebra: Group-Structured Tensors that Enforce Physical Selection Rules

Symmetry determines which physical responses of molecules and crystals are allowed and underlies many models used throughout the physical sciences. Yet several widely used machine-learned models learn symmetry only approximately from data, raising the question of whether approximate symmetry is sufficient for scientific prediction. Using the Orb-v3 force field, we show that it is often not: the relaxed space group of a crystal depends on the orientation of the input for

47\%

of a

150

-crystal subsample, and relaxation lowers the space group for half of all

399

crystals studied. To characterize, quantify, and correct such failures, we introduce a tensor algebra,

\star_G

, defined by convolution over a finite symmetry group and applicable to tensors with noncommuting entries. We prove that the symmetry group determines the algebra's block-diagonal transform and use this structure to analyze model errors on symmetric targets. The resulting projections decompose the average squared error into orthogonal components arising from selection-rule violations, inconsistencies across symmetry-related inputs, and residual error. The removable components reveal how much error arises from broken symmetry, provide reference-free lower bounds on prediction error, and can be eliminated without sacrificing accuracy. For operators, projection followed by blockwise truncation yields the optimal symmetric low-rank approximation. Applying these projections to force predictions reduces both rates of space-group change to at most

2\%

. Exact selection rules, reference-free error bounds, and symmetry-preserving corrections are therefore available for any trained model without retraining.

Paulina Hoyos, Shashanka Ubaru, Dongsung Huh +5

UT Austin · IBM Research · Independent +2

Scientific Discovery Symmetry Discovery Equivariant Neural Networks Tensor Decomposition Equivariant Representation Learning Singular Value Decomposition

Ask PDF

Abstract

Uncertainty Quantification

Explore similar work

Representation Learning for Equivariant Inference with Guarantees

Equivariant Neural Belief Propagation

Exact Symmetry as Algebra: Group-Structured Tensors that Enforce Physical Selection Rules