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

Representation Learning for Equivariant Inference with Guarantees
May 26, 2025Daniel Ordoñez-Apraez, Vladimir Kostić, Alek Fröhlich +3Disentangled Representation Learning Representation Learning
CSML, Italian Institute of Technology, Genova, Italy · Universit`a di Genova, Genova, Italy · University of Novi Sad, Novi Sad, Serbia +2
Equivariant Neural Belief Propagation
Jun 4, 2026Zehua Cheng, Wei Dai, Jiahao SunEquivariant Representation Learning Factor Graph Optimization
Department of Computer Science, University of Oxford, Oxford, United Kingdom · FLock.io, London, United Kingdom
Exact Symmetry as Algebra: Group-Structured Tensors that Enforce Physical Selection Rules
May 19, 2026Paulina Hoyos, Shashanka Ubaru, Dongsung Huh +5Scientific Discovery Symmetry Discovery
UT Austin · IBM Research · Independent +2

May 26, 2025J/K move · Enter open · S save

Representation Learning for Equivariant Inference with Guarantees

Daniel Ordoñez-Apraez, Vladimir Kostić, Alek Fröhlich, Vivien Brandt+2

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

Open paper AskPDF arXiv

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.

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