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

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.Source: §1 · p. 1

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.Source: §1 · p. 1
• 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.Source: §1 · p. 1

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.Source: §1 · p. 1

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.Source: §1 · p. 2

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.Source: §2 · p. 2

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.Source: §2 · p. 2

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.Source: §2 · p. 2
• 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.Source: §2 · p. 2

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.Source: §3.1 · p. 2

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.Source: §3.1 · pp. 2–3

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.Source: Figure 1 · p. 3, §3.5 · p. 4

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.Source: §3.2 · p. 3

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.Source: §3.3 · p. 3

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.Source: §3.4 · p. 4

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.Source: §3.5 · p. 5

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.Source: §4 · p. 5

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.Source: §4.1 · p. 6
• 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.Source: §4.1 · p. 6, Figure 3 · p. 6

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.Source: §4.3 · p. 6

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.Source: §4.4 · p. 7, §4.3 · p. 7

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.Source: §4.4 · p. 8

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.Source: §4.5 · p. 9

• The primary validated setting is E(3)-equivariant UQ for symmetric rank-2 tensors.Source: §5 · p. 9

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.Source: §5 · p. 9

• 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.Source: §5 · p. 9

• 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.Source: Impact Statement · p. 9
• Formal results establish the representation decomposition, matrix-exponential construction, and numerical stability of the loss formulation.Source: A. Theoretical Proofs and Derivations · p. 12
• 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.Source: §A.1 · p. 12
• 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).Source: §A.2 · p. 12
• 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.Source: §A.3 · p. 13, Then the matrix exponential · p. 13
• Spectral functions such as eigenvalue clamping and anisotropic jitter preserve O(3) equivariance.Source: Then the matrix exponential · p. 13

• The implementation compares Gaussian NLL with a Multivariate Laplace NLL for tensor-valued uncertainty estimation.Source: §A.4 · p. 13
• 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.Source: §A.4 · pp. 13–14
• The Laplace NLL is linear in the Mahalanobis distance rather than quadratic, making large residuals contribute less aggressively than under Gaussian NLL.Source: §A.4 · p. 14

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.Source: §A.4 · p. 14

Gradient Stability

• The unmodified Laplace gradient is not uniformly bounded near the SPD boundary, so implementation stability combines eigenvalue clamping with log-tail robustification.Source: §A.4 · p. 14

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.Source: §B.1 · p. 15

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.Source: §B.3 · p. 16

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.Source: §B.3 · p. 16
• 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.Source: §B.3 · pp. 16–17

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.Source: §B.3 · pp. 16–17

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.Source: §B.3 · p. 17

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.Source: §B.3 · pp. 17–18

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.Source: §C.1 · p. 18

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.Source: §C.1 · p. 18

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.Source: §D.1 · p. 19
• 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.Source: §D.1 · p. 19

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%.Source: §D.2 · p. 19

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.Source: §D.3 · p. 19, Table 7 · p. 19

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.Source: §D.4 · pp. 19–20, Table 8 · p. 20

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.Source: §D.5 · p. 20
• 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.Source: §D.5 · p. 20

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.Source: §E.1 · p. 20

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.Source: §E.2 · p. 21

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.Source: §E.2 · p. 21

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.Source: §E.3 · pp. 21–22

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.Source: §E.4 · p. 22, Figure 8 · p. 22

Conditioning

• The condition-number distribution remains numerically well-conditioned for the final model.The supplied passage does not provide a separate numerical sharpness statistic.Source: §E.4 · p. 22, Figure 8 · p. 22

• 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.Source: §E.5 · p. 22
• 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.Source: §E.5 · p. 22

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.Source: §E.6 · p. 22
• 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.Source: §E.6 · p. 22
• More efficient equivariant basis constructions are needed to reduce the overhead of extending uncertainty modeling to complex properties such as full elasticity tensors.Source: §E.6 · pp. 22–23

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.Source: §E.6 · p. 23