cs.LGJun 4, 2026

Equivariant Neural Belief Propagation

Authors: Zehua ChengWei DaiJiahao Sun

Organizations: Department of Computer Science, University of Oxford, Oxford, United Kingdom · FLock.io, London, United Kingdom

Abstract

Probabilistic inference over spatially embedded variables requires beliefs that respect SE(3)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)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)SE(3). On GEOM-QM9 and GEOM-Drugs, ENBP achieves 98.9% conformational coverage at 0.090 A˚\mathring{A} error with sub-second latency -- over 100×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 (107{\sim}10^{-7} vs.\ 10110^{-1} for augmented baselines).

Explore similar work

CardsList
  1. Representation Learning for Equivariant Inference with Guarantees

    May 26, 2025Daniel Ordoñez-Apraez, Vladimir Kostić, Alek Fröhlich +3EquivarianceGeometric Deep Learning