Spring 2024 GRASP SFI: Erik Bekkers, University of Amsterdam, “Fast, Expressive SE(n) Equivariant Networks through Weight-Sharing in Position-Orientation Space”

This will be a hybrid event with in-person attendance in Levine 307 and virtual attendance on Zoom. This week’s speaker will be virtual. There will be an extended Q&A/discussion after the seminar from 4:00 PM to 4:30 PM. 

ABSTRACT

Based on the theory of homogeneous spaces we derive \textit{geometrically optimal edge attributes} to be used within the flexible message passing framework. We formalize the notion of weight sharing in convolutional networks as the sharing of message functions over point-pairs that should be treated equally. We define equivalence classes of point-pairs that are identical up to a transformation in the group and derive attributes that uniquely identify these classes. Weight sharing is then obtained by conditioning message functions on these attributes. As an application of the theory, we develop an efficient equivariant group convolutional network for processing 3D point clouds. The theory of homogeneous spaces tells us how to do group convolutions with feature maps over the homogeneous space of positions ℝ3, position and orientations ℝ3×S2, and the group SE(3) itself. Among these, ℝ3×S2 is an optimal choice due to the ability to represent directional information, which ℝ3 methods cannot, and it significantly enhances computational efficiency compared to indexing features on the full SE(3) group. We empirically support this claim by reaching state-of-the-art results — in accuracy and speed — on three different benchmarks: interatomic potential energy prediction, trajectory forecasting in N-body systems, and generating molecules via equivariant diffusion models.