PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2608.24386v1cs.LGcs.AI

TL;DR

Tensor-valued E(3)-equivariant predictors generally lack rigorous confidence measures, while covariance parameterizations must satisfy both equivariance and positive-definiteness. The paper addresses this with an irreducible-representation covariance head, matrix-exponential SPD mapping, and robust Log-Euclidean scoring, achieving competitive validation results with symmetry-preserving uncertainty estimates.

  • Problem

    E(3)-equivariant tensor predictors are deterministic, while valid uncertainty requires a 6 × 6 covariance that transforms equivariantly in Kelvin-Mandel coordinates.

  • Method

    The framework decomposes the covariance into irreducible representations, predicts a symmetric operator, maps it to SPD covariance by matrix exponentiation, and trains with LE-ESO.

  • Results

    The method achieves competitive performance on inertia and dielectric tensor prediction, including MAE 1.55 on Materials Project dielectric prediction and 10^-7-level equivariance errors.

  • Takeaways & Limitations

    Full-covariance, symmetry-preserving uncertainty remains competitive for prediction and can identify anisotropic tensor failure modes beyond diagonal or trace-based approximations.

  • Takeaways & Limitations

    Higher-order empirical validation is not exhaustive, and broader testing across equivariant backbones, tensor orders, and symmetry groups remains future work.

Abstract

from arXiv · show

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.

1. Introduction

Tensor-valued predictors need uncertainty estimates that respect the tensor’s rotational symmetry, but existing approaches struggle to guarantee both equivariance and valid covariance structure. The paper introduces an equivariant full-covariance framework using irreducible representations, matrix exponentiation, and Log-Euclidean scoring.

  • Motivation: Existing E(3)-equivariant networks provide deterministic tensor predictions without confidence measures, risking overconfident scientific decisions.Applications include material properties, biomedical imaging, and computational fluid dynamics.
  • Problem formulation: Symmetric rank-2 tensors have six independent Kelvin-Mandel coordinates, so their full uncertainty requires a 6 × 6 covariance matrix.Rotations act through an orthogonal six-dimensional representation on these coordinates.
  • Problem formulation: The covariance must transform equivariantly under rotations, creating a parameterization challenge for uncertainty-aware tensor prediction.The required covariance transformation law is incompatible with simply combining standard SPD and direct equivariant parameterizations.
  • Contributions: The proposed framework uses an equivariant matrix-exponential head based on Sym^2(ρc) ≅ 2 × (ℓ=0) ⊕ 2 × (ℓ=2) ⊕ 1 × (ℓ=4).It also introduces a Log-Euclidean scoring objective for joint uncertainty and feature optimization.
  • Validation: MAE 1.55 is reported for Materials Project dielectric prediction, alongside validation on ModelNet40 inertia tensors.The framework is presented as providing competitive performance with symmetry-preserving uncertainty estimates.

2. Related Work

Prior work provides strong equivariant tensor point prediction and several SPD covariance parameterizations, but does not jointly deliver full tensor-valued equivariant uncertainty. Existing probabilistic extensions commonly rely on heuristic, diagonal, isotropic, or computationally limited approaches.

  • Equivariant Tensor Prediction: Existing E(3)-equivariant networks achieve strong tensor prediction but are deterministic and do not quantify confidence.The paper positions its framework as addressing this uncertainty gap.
  • SPD Constraints in Neural Networks: Cholesky, eigendecomposition, and Riemannian methods guarantee SPD in fixed coordinates but do not by themselves ensure equivariant covariance maps.The challenge concerns the interaction between parameterization and the covariance representation ρc.
  • Probabilistic Methods: Ensembles offer heuristic uncertainty estimates, while the proposed framework learns the full 21-parameter aleatoric covariance tensor for anisotropic noise.The cited motivation is modeling inherent anisotropic noise in physical properties.
  • Probabilistic Methods: Equivariant Bayesian approaches and Gaussian processes often target scalar or vector quantities or restrict tensor uncertainty to isotropic or diagonal covariances.Scaling to the full covariance structure of rank-2 tensors remains difficult in the related methods described.

3. Methods

The method predicts symmetric rank-2 tensor distributions in Kelvin-Mandel coordinates, enforcing O(3) equivariance through representation-specific operator construction and enforcing SPD through matrix exponentiation. A Log-Euclidean objective trains the mean and covariance jointly while supporting robust optimization.

  • Problem Formulation: The framework predicts a full distribution for symmetric rank-2 tensors, using a 6 × 6 covariance over Kelvin-Mandel coordinates.It applies to 3D point clouds with rank-2 tensor attributes, including dielectric tensors.
  • Equivariance: Both mean and covariance are translation-invariant for global properties and equivariant under orthogonal transformations.The construction supports rotations and reflections through the even-parity Kelvin-Mandel representation.
  • Positive-Definite Mapping: The matrix exponential maps the flat space sym(6) to the SPD manifold while preserving equivariance under orthogonal conjugation.The network therefore predicts unconstrained symmetric operators before producing valid covariances.
  • Kelvin-Mandel Representation: Kelvin-Mandel notation preserves the tensor-vector isometry and yields covariance transformation by orthogonal conjugation.This maintains coordinate invariance of the physical uncertainty.
  • Representation Decomposition: The symmetric covariance representation decomposes as Sym^2(ρc) ≅ 2 × (ℓ=0) ⊕ 2 × (ℓ=2) ⊕ 1 × (ℓ=4), giving 21 independent degrees of freedom.The ℓ=1 and ℓ=3 terms belong to the antisymmetric component and are excluded by covariance symmetry.
  • Covariance Construction: The covariance head assembles an equivariant symmetric operator A(X) from structured irreducible features, then sets Σ(X) = exp(A(X)).Fixed Clebsch-Gordan basis matrices and equivariant coefficients provide hard-constrained geometric consistency.
  • Training Objective: LE-ESO combines uncertainty-volume regularization with a robustified data-fit term, and α = 1 recovers the multivariate Laplace negative log-likelihood in Log-Euclidean form.The objective trades strict propriety for outlier stability when robustified or when α ≠ 1; α = 1 gave the lowest validation MAE in the reported sweep.

4. Experiments

Experiments evaluate the framework on controlled ModelNet40 geometry and Materials Project dielectric tensors, testing equivariance, SPD validity, accuracy, calibration, risk utility, and OOD sensitivity. Results show near-machine-precision symmetry preservation, competitive accuracy, calibrated uncertainty, and useful directional and shift-sensitive risk signals.

  • Experimental setup: Experiments cover ModelNet40 inertia tensors and Materials Project dielectric tensors, with additional studies of shape covariance, rank-4 elasticity, runtime, and sensitivity analyses.The main comparisons use deterministic MSE and diagonal-UQ baselines, while ablations include non-equivariant and non-SPD variants.
  • Controlled geometric validation: Equivariance errors remain at 10^-7, while the full-covariance model reduces MAE by 15% relative to the diagonal baseline and maintains >99.9% SPD validity.The median covariance condition number is 6.8.
  • Controlled geometric validation: Uncertainty ellipsoids align with principal shape axes, expand in sparse-density regions, and preserve tensorial correlations across components.ModelNet40 visualizations report sample errors of 0.293–0.472 and Mahalanobis distances of 0.37–0.50.
  • Prediction accuracy: Materials Project MAE is 1.55, outperforming MACE deep ensemble at 1.96 and diagonal UQ at 2.25 while remaining near deterministic point-prediction baselines.The reported deterministic MAEs are 1.91 for DTNet and 1.41 for GoeCTP.
  • Uncertainty calibration: The Materials Project reliability diagram reports MACE 0.0489, with slightly below-diagonal high-confidence behavior indicating conservative under-confidence.The reliability assessment uses the same Multivariate Laplace distribution as the training objective.
  • Risk and OOD sensitivity: At 90% coverage, λmax(Σ) ranking improves retained-set MAE by 3.1% versus 0.8% for Trace(Σ), and predicted λmax rises from 2.71 to 5.23 under complete chemical substitution.The directional-ranking advantage over Trace is smaller but consistent, while chemical-shift analysis does not disentangle aleatoric and epistemic uncertainty.
  • Ablation and validity: The matrix-exponential head is the only tested head achieving O(10^-7) equivariance error with strict SPD validity, whereas Cholesky shows EΣ ≈ 0.43.The Cholesky failure is attributed to its lower-triangular structure not being preserved under orthogonal conjugation.

5. Discussion

The framework is primarily validated for E(3)-equivariant uncertainty quantification of symmetric rank-2 tensors, with limited but promising evidence for higher-order targets. Broader empirical validation remains future work.

  • The primary validated setting is E(3)-equivariant UQ for symmetric rank-2 tensors.
  • Extensibility to Higher-Order Tensors: A rank-4 elasticity experiment provides supporting evidence, rather than a comprehensive benchmark, for extending the construction beyond rank-2 tensors.The experiment reports competitive MAE, improved empirical coverage, and preserved numerical equivariance and SPD validity.
  • Broader validation across equivariant backbones, tensor orders, and symmetry groups remains future work.Extension to other groups depends on suitable representation-theoretic bases and implementation tools.

Impact Statement

The paper’s potential impact is in scientific machine learning applications involving tensor-valued material properties, supported by formal guarantees for equivariant covariance construction. These guarantees rely on a 21-degree-of-freedom covariance representation and matrix-exponential mapping.

  • The framework targets scientific machine learning tasks involving tensor-valued material properties such as dielectric and elastic responses.The stated potential impact includes materials discovery tasks where these properties are important.
  • Formal results establish the representation decomposition, matrix-exponential construction, and numerical stability of the loss formulation.
  • Kelvin-Mandel transformations are constructed from the Kronecker product and projection onto the six-dimensional symmetric subspace.The selection matrix and scaling matrix define the transformation, which preserves orthogonality in Kelvin-Mandel space.
  • The covariance representation has 21 independent degrees of freedom, matching a symmetric 6 × 6 covariance matrix.The decomposition is 2 × (l = 0) ⊕ 2 × (l = 2) ⊕ 1 × (l = 4).
  • The matrix exponential maps a symmetric Lie-algebra output to an SPD covariance while preserving equivariance under orthogonal conjugation.Positive eigenvalues ensure strict positive-definiteness, and the exponential commutes with similarity transformations.
  • Spectral functions such as eigenvalue clamping and anisotropic jitter preserve O(3) equivariance.

A.4. Numerically Stable Loss Function

The loss is formulated in Log-Euclidean coordinates for Gaussian and Multivariate Laplace likelihoods, with the Laplace form using a linear Mahalanobis penalty for greater robustness. A Huber-style logarithmic tail further limits extreme-outlier influence, while the architecture jointly predicts mean and covariance.

  • The implementation compares Gaussian NLL with a Multivariate Laplace NLL for tensor-valued uncertainty estimation.
  • Log-Euclidean likelihoods replace log det Σ with Tr(A) and use exp(−A) for the inverse covariance when Σ = exp(A).These identities yield numerically stable forms in the symmetric Lie algebra.
  • The Laplace NLL is linear in the Mahalanobis distance rather than quadratic, making large residuals contribute less aggressively than under Gaussian NLL.
  • Robustification: The Laplacian-Huber loss is linear below τ and logarithmic above τ, with τ = 5.0 selected so approximately 99% of well-predicted samples remain below the threshold.
  • Gradient Stability: The unmodified Laplace gradient is not uniformly bounded near the SPD boundary, so implementation stability combines eigenvalue clamping with log-tail robustification.
  • Architecture: The equivariant architecture jointly predicts the mean and covariance through separate heads, with the covariance head mapping irreducible features to a symmetric Lie-algebra matrix.The dual-head design preserves the underlying geometric symmetries for both outputs.

B.3. Training Protocol and Stability Measures

The framework combines staged optimization, bounded spectral parameterization, and robustification to stabilize equivariant covariance learning while preserving SPD and symmetry guarantees.

  • Optimization Strategy: Training transitions from a five-epoch MSE warmup to full LE-ESO optimization to establish reliable mean predictions before uncertainty learning.The auxiliary MSE weight decays from 0.9 to zero by epoch 5.
  • Numerical Stabilization: Anisotropic jitter separates degenerate covariance eigenvalues, preventing singular eigenvalue-decomposition Jacobians during differentiation.The perturbation uses ε ≈ 10^-6 and preserves equivariance errors near 10^-7.
  • Numerical Stabilization: Eigenvalue clamping to [−4, 3] bounds covariance eigenvalues in approximately [0.018, 20.1], preventing collapse and explosion during early training.The bounds are applied before matrix exponentiation.
  • Geometric Parameterization: Working in the flat symmetric-matrix Lie algebra and mapping with the matrix exponential provides geometrically defined optimization while preserving positive-definite covariances.The exponential of the symmetric parameter matrix remains SPD, and the loss preserves O(3)-invariance.
  • Robust Objective: The Laplacian-Huber loss applies a linear penalty below Mahalanobis distance 5.0 and a logarithmic tail penalty above it to reduce extreme-outlier influence.The log-tail gradient approaches a finite constant in a dominant direction, mitigating rare extreme residuals.
  • Symmetry Preservation: The objective is an exact symmetry-preserving training loss rather than an approximate regularization or augmentation scheme.Trace and Mahalanobis terms remain invariant under orthogonal transformations.

C.1. Dataset Configuration and Preprocessing

Experiments use complementary geometric and materials datasets, controlled baselines, and runtime and hyperparameter analyses to assess structured equivariant SPD uncertainty prediction.

  • ModelNet40: ModelNet40 provides controlled geometric validation using 12,311 CAD models, official train/test splits, 2,048 surface points, and Gaussian jitter with σnoise = 0.01.The injected noise supplies aleatoric uncertainty for testing geometric ambiguity.
  • Materials Project: The Materials Project experiment extends evaluation to real-data dielectric tensors and filters structures by size and tensor positive-definiteness.The supplied passage states exclusions below 3 or above 30 atoms and verifies strictly positive tensor eigenvalues.
  • ModelNet40: The shape-covariance benchmark tests whether the construction generalizes beyond inertia tensors, while acknowledging that a closed-form estimator exists.It is presented as controlled validation rather than a setting requiring neural prediction.
  • ModelNet40: The UQ model matches the deterministic baseline in point-prediction MAE/RMSE while producing exactly E(3)-equivariant, strictly SPD predictive covariances.Equivariance errors are on the order of 10^-7.
  • Higher-Order Tensor Validation: On rank-4 elasticity prediction, test MAE is approximately 5.0 GPa, while coverage rises from approximately 35% to 52% and uncertainty–error correlation from −0.15 to 0.31.Both numerical consistency and predictive-covariance SPD validity remain at 100%.
  • Efficiency: A single forward pass produces the full anisotropic covariance, while diagonal-UQ overhead is 1.5%; full-covariance cost is driven mainly by the covariance branch and backpropagation.Runtime profiling uses an RTX 4060 Ti with batch size 16.
  • Sensitivity Analysis: Across α values from 0.03 to 1.00, validation MAE remains between 0.352 and 0.457, with α = 1 achieving the lowest value.LE-ESO values are not directly comparable across α because the objective changes.

D.5. Additional Risk-Coverage Analysis

Risk-coverage analysis shows that directional full-covariance uncertainty improves retained-set selection, while training diagnostics confirm stable optimization and structural guarantees.

  • Risk-Coverage Ranking: At 90% coverage, λmax(Σ) ranking improves retained-set MAE by 3.1% relative to the full test set.Its advantage over Trace(Σ) is approximately 1.5%, indicating a moderate rather than dramatic difference.
  • Risk-Coverage Ranking: The diagonal-UQ baseline ranks test examples less informatively than λmax or Trace rankings from the full-covariance model.These comparisons support the role of off-diagonal correlations and directional uncertainty.
  • Training Dynamics: The five-epoch warmup transitions to heavy-tailed LE-ESO without numerical spikes, while point-prediction MAE remains stable after adding the uncertainty branch.The training dynamics also track the balance between Mahalanobis fit and log-determinant regularization.
  • Structural Verification: Equivariance error remains at 10^-7 across tested rotations, and predicted covariance minimum eigenvalues remain strictly positive above 10^-5.These checks verify near-machine-precision equivariance and no observed SPD violations during training.
  • Numerical Conditioning: Materials Project covariance condition numbers have mean 3.80 and maximum 16.4, indicating numerically well-conditioned predictions without auxiliary spectral regularization.The analysis concerns dielectric-tensor uncertainty arising from physical and DFT approximation errors.
  • O(3) Equivariance: Improper-rotation tests also achieve errors at 10^-7, extending symmetry validation from SO(3) to full O(3).This supports validity for coordinate systems with either handedness.

E.4. Spectral Analysis and Sharpness Distribution

Spectral analysis reports strictly positive covariance spectra and well-conditioned uncertainty estimates, supporting manifold consistency and numerical stability.

  • Eigenvalue Spectrum: The covariance eigenvalue distribution has minimum eigenvalue λmin ≈ 0.449, confirming strict positive-definiteness and avoidance of variance collapse.Figure 8 is used for spectral validation.
  • Conditioning: The condition-number distribution remains numerically well-conditioned for the final model.The supplied passage does not provide a separate numerical sharpness statistic.

E.5. ModelNet40 SPD Analysis

The framework produces physically meaningful, symmetry-consistent uncertainty estimates for inertia tensors, with strict SPD validity and well-conditioned covariance structures.

  • Uncertainty ellipsoids align with principal shape axes, expand where point density is sparse, and preserve tensorial correlations across components.These patterns demonstrate E(3)-equivariance and capture sampling ambiguity.
  • More than 99.9% of predicted covariances satisfy strict SPD requirements, while the median condition number is 6.8.The analysis contrasts these results with unconstrained baselines that frequently violate physical constraints.

E.6. Limitations and Future Work

The framework is modular and backbone-agnostic, but extending it to higher-order tensors faces substantial basis-construction and covariance-size scaling challenges.

  • Limitations and Future Work: Extending full parameterization to higher-order tensors requires group- and representation-specific basis construction, with rank-4 evidence remaining preliminary.The main implementation and empirical validation focus on symmetric rank-2 tensors.
  • Limitations and Future Work: For rank-4 and higher tensors, basis enumeration scales as O(d^ℓ), while covariance dimensions grow combinatorially to O(k^4).These costs may require low-rank covariance factorization, hierarchical uncertainty modeling, and careful memory management.
  • Limitations and Future Work: More efficient equivariant basis constructions are needed to reduce the overhead of extending uncertainty modeling to complex properties such as full elasticity tensors.
  • Modularity and Backbone Extensibility: The matrix-exponential UQ head can integrate with any E(3)-equivariant architecture without architectural reengineering.The study uses a standard message-passing backbone, while future work considers higher-accuracy architectures such as GoeCTP.
Loading 2608.24386v1…