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

Motivation

• Known: E(3)-equivariant neural network(ENN)은 dielectric tensor 등 tensor 속성 예측에서 뛰어난 점 추정(point estimate) 성능을 보여왔으며, Cholesky decomposition·eigendecomposition·Riemannian optimization 등을 통해 SPD 공분산을 보장하는 기법들도 알려져 있다.
• Gap: 기존 Cholesky류 covariance head는 SPD를 보장하지만 Kelvin-Mandel 공분산 좌표에서 equivariant하지 않고, 반대로 직접적인 equivariant regression은 회전 변환 법칙은 보존하지만 SPD를 보장하지 못한다. 즉 equivariance와 SPD 제약을 동시에 만족시키는 파라미터화가 부재했으며, 기존 equivariant Bayesian/GP 접근은 scalar나 vector, 혹은 등방성/대각 공분산에 국한되어 rank-2 tensor의 완전한 21-parameter 공분산 구조를 다루지 못했다.
• Why: 과학적 의사결정(재료 설계, 의료 영상 등)에서 tensor 예측의 과신(overconfidence)은 값비싼 실험 실패로 이어질 수 있으므로, 물리적으로 일관되고 대칭성을 보존하는 신뢰도(confidence) 측정치를 제공하는 것이 중요하다.
• Approach: 공분산을 Lie algebra sym(6) 상의 무제약 대칭 연산자 A(X)로 예측한 뒤 matrix exponentiation을 통해 SPD manifold P6로 사상함으로써 equivariance와 SPD를 동시에 엄격히 보장하고, 이를 Multivariate Laplace 분포 기반의 강건한 surrogate loss인 Log-Euclidean Equivariant Scoring Objective(LE-ESO)로 학습한다.

Achievement

Figure 4. Performance on Materials Project. (a) Parity plot of diagonal components (R2 = 0.659). (b) Multivariate Laplac

1. Equivariant SPD parameterization 해결: Sym2(ρc) ≅ 2×(l=0) ⊕ 2×(l=2) ⊕ 1×(l=4) 기약표현 분해에 기반한 equivariant matrix-exponential head를 제안하여, Σ(X) = exp(A(X))가 exp(ρc(R)Aρc(R)⊤) = ρc(R)exp(A)ρc(R)⊤ 관계를 만족함으로써 모든 회전 프레임에서 동시에 SPD와 equivariance를 보장함을 이론적으로 증명했다.
2. 안정적 결합 최적화 목적함수 제안: 불확실성 보정(calibration)과 기하학적 특징 추출을 통합하는 LE-ESO를 도입하여 이상치(heavy-tailed error)에 강건하고 안정적인 학습을 가능하게 했다.
3. 실증적 검증: ModelNet40 inertia tensor와 Materials Project dielectric tensor 데이터셋에서 검증하여, MAE 1.55의 경쟁력 있는 예측 성능과 함께 물리적으로 일관되고 대칭성을 보존하는 불확실성 추정치, 유용한 risk 및 OOD(out-of-distribution) 민감도를 보여주었다.

How

Figure 1. E(3)-equivariant tensor uncertainty framework. (1) Feature Decomposition: covariance head models symmetry via

• Kelvin-Mandel 표기법을 사용해 대칭 rank-2 tensor C를 6차원 벡터 cKM으로 평탄화하며, 이는 Frobenius norm을 보존하는 isometry 성질을 가짐 (∥cKM∥2 = ∥C∥F).
• 회전 R 하에서 c′KM = ρc(R)cKM이고 공분산은 Σ′ = ρc(R)Σρc(R)⊤로 변환됨을 만족하도록 설계.
• E(3)-equivariant encoder(L layers)로부터 mean head는 l=0⊕l=2 feature로 μ(X)를 예측하고, covariance head는 Sym2(ρc) basis (2×l=0 ⊕ 2×l=2 ⊕ 1×l=4)로 A(X) ∈ sym(6)를 예측.
• A(X)를 matrix exponentiation을 통해 SPD manifold P6로 사상하여 Σ(X) = exp(A(X))를 얻음 (Log-Euclidean 방식).
• 학습 목적함수는 Multivariate Laplace negative log-likelihood를 일반화한 LE-ESO를 사용해 이상치에 강건하게 mean과 covariance를 결합 최적화.
• e3nn 라이브러리의 spherical harmonic basis를 활용해 equivariant network 구현.
• O(3) 전체(반사 포함) equivariance는 even-parity Kelvin-Mandel basis ρc(R) = R⊗sR로 처리하며, Section 4.5(Table 4) 및 Appendix E.3에서 수치적으로 검증.

Originality

• Cholesky/eigendecomposition 기반 SPD 보장 기법과 direct equivariant regression 간의 근본적 트레이드오프(SPD vs equivariance)를 명확히 규명하고, matrix exponential을 통한 Lie algebra→SPD manifold 사상으로 두 성질을 동시에 만족시키는 최초의 해법을 제시.
• 대칭 rank-2 tensor 공분산을 Sym2(ρc) ≅ 2×(l=0) ⊕ 2×(l=2) ⊕ 1×(l=4)로 기약표현 분해하여 equivariant covariance head를 설계한 것은 새로운 접근.
• 기존 equivariant Bayesian/GP 방법이 scalar·vector 혹은 등방성/대각 공분산에 국한된 것과 달리, 21-parameter 전체 공분산 구조를 신경망으로 확장 가능하게 함.
• Multivariate Laplace 분포와 Log-Euclidean metric을 결합한 LE-ESO라는 새로운 강건 손실함수를 제안.

Limitation & Further Study

• 본 프레임워크는 rank-2 대칭 tensor에 특화되어 있으며, A(X)의 파라미터화는 tensor 차수 및 대칭군마다 별도로 구성해야 하므로 더 높은 rank tensor(예: elasticity tensor)로의 일반화에는 추가 작업이 필요하다.
• matrix exponentiation 연산의 계산 비용 및 수치적 안정성(특히 6×6 이상 큰 행렬에서의 gradient 계산)에 대한 상세 분석이 제한적일 수 있다.
• 검증이 ModelNet40 inertia tensor와 Materials Project dielectric tensor 두 데이터셋에 한정되어 있어, 더 다양한 도메인(예: diffusion MRI, CFD)에서의 일반화 성능 검증이 필요하다.
• 최근 제안된 High-Rank Irreducible Cartesian Tensor(ICT) decomposition과 같은 대안적 equivariant basis와의 비교 및 결합 가능성은 추후 연구로 남아있다.

Evaluation

총평: Equivariance와 SPD 제약이라는 두 근본적으로 상충하는 요구사항을 matrix exponential을 통한 Lie algebra-to-manifold 사상으로 우아하게 해결한, 이론적으로 탄탄하고 실용적 의미가 큰 연구이다. 다만 검증 범위가 두 데이터셋에 제한되어 있어 더 넓은 도메인과 더 높은 rank tensor로의 확장성 검증이 후속 연구로 요구된다.