Equivariant Covariance Tensors: Guaranteed SPD Uncertainty for Tensor-Valued Geometric Learning
cs.LG, cs.AI
Submitted: 2026-08-25
Updated: 2026-08-25
Comments: Accepted to ICML 2026
License: http://creativecommons.org/licenses/by/4.0/
The gist: Tensor-valued prediction is fundamental to geometric deep learning, yet uncertainty quantification (UQ) for such outputs remains an open challenge.
Terminology
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 Sym 2(ρ c) 2 times(l=0) 2 times(l=2) 1 times(l=4). By mapping from the flat Lie algebra 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.
Sources
- Open Materials 2024 (OMat24) Inorganic Materials Dataset and Models
- Emergent Equivariance in Deep Ensembles
- Equivariant Graph Neural Networks for Prediction of Tensor Material Properties of Crystals
- Neural Network Layers for Prediction of Positive Definite Elastic Stiffness Tensors
- Multi-Scale Message Passing Neural PDE Solvers
- A Structured Review of Literature on Uncertainty in Machine Learning & Deep Learning
- e3nn: Euclidean Neural Networks
- BI-EqNO: Generalized Approximate Bayesian Inference with an Equivariant Neural Operator Framework
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks