Equivariant Neural Belief Propagation

2026-06-04

AI Summary

This work addresses the challenge of modeling anisotropic uncertainty and multimodal distributions under SE(3) symmetry, which existing methods struggle to handle effectively. The authors propose the first SE(3)-equivariant neural belief propagation framework, wherein messages are represented as equivariant Gaussian mixture models. Rank-2 precision matrices are generated via equivariant outer products, and differentiable spectral decomposition combined with KL-based greedy mixture reduction preserves strict equivariance while preventing mode collapse. Evaluated on GEOM-QM9 and GEOM-Drugs, the method achieves a conformation coverage of 98.9% with a mean error of 0.090 Å, offering inference speeds over 100× faster than diffusion models. In multi-agent robotic inference tasks, it yields near-zero collision rates and attains equivariance errors on the order of 10⁻⁷.

Abstract

Probabilistic inference over spatially embedded variables requires beliefs that respect $SE(3)$ symmetry, yet existing equivariant networks produce only scalars and vectors -- not the rank-2 precision tensors needed for anisotropic uncertainty, and single-component messages collapse multi-modal energy landscapes to physically meaningless averages. We introduce Equivariant Neural Belief Propagation (ENBP), a factor-graph framework whose messages are equivariant Gaussian mixture models with sufficient statistics that transform exactly under $SE(3)$. Rank-2 precision matrices are synthesised via equivariant outer products, ingested through differentiable spectral decomposition, and kept tractable by a greedy KL-based mixture reduction that provably commutes with $SE(3)$. On GEOM-QM9 and GEOM-Drugs, ENBP achieves 98.9% conformational coverage at 0.090 $\mathring{A}$ error with sub-second latency -- over $100\times$ faster than diffusion baselines at higher accuracy. On multi-body robotic inference, vanilla loopy BP diverges at 15+ agents while ENBP converges with near-zero collision rates and machine-precision equivariance error (${\sim}10^{-7}$ vs.\ $10^{-1}$ for augmented baselines).

Problem

• Equivariance
• SE(3) symmetry
• Probabilistic inference
• Anisotropic uncertainty
• Belief propagation

Innovation

• Equivariant Neural Belief Propagation
• SE(3) equivariance
• Gaussian mixture models
• rank-2 precision tensors
• differentiable spectral decomposition