Many machine learning tasks involve learning functions that are known to be invariant or equivariant to certain symmetries of the input data. However, it is often challenging to design neural network architectures that respect these symmetries while being expressive and computationally efficient. For example, Euclidean motion invariant/equivariant graph or point cloud neural networks. We introduce Frame Averaging (FA), a general purpose and systematic framework for adapting known (backbone) architectures to become invariant or equivariant to new symmetry types. Our framework builds on the well known group averaging operator that guarantees invariance or equivariance but is intractable. In contrast, we observe that for many important classes of symmetries, this operator can be replaced with an averaging operator over a small subset of the group elements, called a frame. We show that averaging over a frame guarantees exact invariance or equivariance while often being much simpler to compute than averaging over the entire group. Furthermore, we prove that FA-based models have maximal expressive power in a broad setting and in general preserve the expressive power of their backbone architectures. Using frame averaging, we propose a new class of universal Graph Neural Networks (GNNs), universal Euclidean motion invariant point cloud networks, and Euclidean motion invariant Message Passing (MP) GNNs. We demonstrate the practical effectiveness of FA on several applications including point cloud normal estimation, beyond $2$-WL graph separation, and $n$-body dynamics prediction, achieving state-of-the-art results in all of these benchmarks.
FAENet: Frame Averaging Equivariant GNN for Materials Modeling
Mila – Quebec AI Institute（米拉-魁北克人工智能研究所）；Université Paris-Saclay（巴黎萨克雷大学）；CentraleSupélec（中央理工-高等电力学院）；Inria（法国国家信息与自动化研究所）；Intel Labs（英特尔实验室）；Université de Montréal（蒙特利尔大学）；McGill Unversity（麦吉尔大学）
It is widely believed that engineering a model to be invariant/equivariant improves generalisation. Despite the growing popularity of this approach, a precise characterisation of the generalisation benefit is lacking. By considering the simplest case of linear models, this paper provides the first provably non-zero improvement in generalisation for invariant/equivariant models when the target distribution is invariant/equivariant with respect to a compact group. Moreover, our work reveals an interesting relationship between generalisation, the number of training examples and properties of the …
Group Averaging for Physics Applications: Accuracy Improvements at Zero Training Cost
Center for Cosmology and Particle Physics（宇宙与粒子物理中心）；Department of Physics（物理系）；New York University（纽约大学）；Department of Applied Mathematics & Statistics and Mathematical Institute for Data Science（应用数学与统计系及数据科学数学研究所）；Johns Hopkins University（约翰霍普金斯大学）
NVIDIA（英伟达）；Weizmann Institute of Science（魏茨曼科学研究所）；University of Toronto（多伦多大学）；Vector Institute（向量研究所）
The task of shape space learning involves mapping a train set of shapes to and from a latent representation space with good generalization properties. Often, real-world collections of shapes have symmetries, which can be defined as transformations that do not change the essence of the shape. A natural way to incorporate symmetries in shape space learning is to ask that the mapping to the shape space (encoder) and mapping from the shape space (decoder) are equivariant to the relevant symmetries. In this paper, we present a framework for incorporating equivariance in encoders and decoders by int…
In this work we develop a new method, named Sub-graph Permutation Equivariant Networks (SPEN), which provides a framework for building graph neural networks that operate on sub-graphs, while using a base update function that is permutation equivariant, that are equivariant to a novel choice of automorphism group. Message passing neural networks have been shown to be limited in their expressive power and recent approaches to over come this either lack scalability or require structural information to be encoded into the feature space. The general framework presented here overcomes the scalabilit…
Chinese Academy of Sciences（中国科学院）；Xi’an Jiaotong University（西安交通大学）；George Mason University（乔治梅森大学）；Microsoft Research（微软研究院）
Group equivariance (e.g. SE(3) equivariance) is a critical physical symmetry in science, from classical and quantum physics to computational biology. It enables robust and accurate prediction under arbitrary reference transformations. In light of this, great efforts have been put on encoding this symmetry into deep neural networks, which has been shown to improve the generalization performance and data efficiency for downstream tasks. Constructing an equivariant neural network generally brings high computational costs to ensure expressiveness. Therefore, how to better trade-off the expressiven…
Equivariant Spatio-Temporal Attentive Graph Networks to Simulate Physical Dynamics
Gaoling School of Artificial Intelligence, Renmin University of China（中国人民大学高瓴人工智能学院）；North Carolina State University（北卡罗来纳州立大学）；Institute for AI Industry Research (AIR), Tsinghua University（清华大学人工智能产业研究院）；Tencent AI Lab（腾讯AI实验室）
Quantitative Approximation Rates for Group Equivariant Learning
Department of Mathematics（数学系）；Texas A&M University（德克萨斯A&M大学）；Faculty of Mathematics（数学系）；Technion- Israel Institute of Technology（技术学院-以色列理工学院）；Computer Science and Artificial Intelligence Laboratory（计算机科学与人工智能实验室）；Massachusetts Institute of Technology（麻省理工学院）
University of Padova（帕多瓦大学）；Stanford University（斯坦福大学）；TU Munich（慕尼黑工业大学）；TU Dortmund（多特蒙德工业大学）；Google（谷歌）
We present a 3D capsule module for processing point clouds that is equivariant to 3D rotations and translations, as well as invariant to permutations of the input points. The operator receives a sparse set of local reference frames, computed from an input point cloud and establishes end-to-end transformation equivariance through a novel dynamic routing procedure on quaternions. Further, we theoretically connect dynamic routing between capsules to the well-known Weiszfeld algorithm, a scheme for solving \emph{iterative re-weighted least squares} (IRLS) problems with provable convergence propert…
Institute of Mathematics（数学研究所）；Image Computing, University of Luebeck, Luebeck, Germany（图像计算，吕贝克大学，吕贝克，德国）；Department of Applied Mathematics（应用数学系）；Theoretical Physics, University of Cambridge, Cambridge, United Kingdom（理论物理，剑桥大学，剑桥，英国）
Purdue University（普渡大学）；Imperial College London（伦敦帝国理工学院）；Twitter（推特）；MIT（麻省理工学院）；UCSD（加州大学圣地亚哥分校）；University of Tuebingen（图宾根大学）；NVIDIA Research（英伟达研究院）
Message-passing neural networks (MPNNs) are the leading architecture for deep learning on graph-structured data, in large part due to their simplicity and scalability. Unfortunately, it was shown that these architectures are limited in their expressive power. This paper proposes a novel framework called Equivariant Subgraph Aggregation Networks (ESAN) to address this issue. Our main observation is that while two graphs may not be distinguishable by an MPNN, they often contain distinguishable subgraphs. Thus, we propose to represent each graph as a set of subgraphs derived by some predefined po…
École Normale Supérieure, PSL（巴黎高等师范学院（PSL大学））；Mila（米拉研究所）；Université de Montréal（蒙特利尔大学）；McGill University（麦吉尔大学）
The Strong Lottery Ticket Hypothesis (SLTH) stipulates the existence of a subnetwork within a sufficiently overparameterized (dense) neural network that -- when initialized randomly and without any training -- achieves the accuracy of a fully trained target network. Recent works by Da Cunha et. al 2022; Burkholz 2022 demonstrate that the SLTH can be extended to translation equivariant networks -- i.e. CNNs -- with the same level of overparametrization as needed for the SLTs in dense networks. However, modern neural networks are capable of incorporating more than just translation symmetry, and …
A General Framework for Equivariant Neural Networks on Reductive Lie Groups
University of Cambridge（剑桥大学）；ENS Paris-Saclay（巴黎萨克雷高等师范学院）；Université Paris-Saclay（巴黎萨克雷大学）；Massachusetts Institute of Technology（麻省理工学院）；EIA University（EIA大学）；University of British Columbia（不列颠哥伦比亚大学）
Chalmers University of Technology（查尔姆斯理工大学）；TU Wien（维也纳工业大学）
Lattice gauge equivariant convolutional neural networks (L-CNNs) are a framework for convolutional neural networks that can be applied to non-Abelian lattice gauge theories without violating gauge symmetry. We demonstrate how L-CNNs can be equipped with global group equivariance. This allows us to extend the formulation to be equivariant not just under translations but under global lattice symmetries such as rotations and reflections. Additionally, we provide a geometric formulation of L-CNNs and show how convolutions in L-CNNs arise as a special case of gauge equivariant neural networks on SU…
We propose a semantic segmentation model that exploits rotation and reflection symmetries. We demonstrate significant gains in sample efficiency due to increased weight sharing, as well as improvements in robustness to symmetry transformations. The group equivariant CNN framework is extended for segmentation by introducing a new equivariant (G->Z2)-convolution that transforms feature maps on a group to planar feature maps. Also, equivariant transposed convolution is formulated for up-sampling in an encoder-decoder network. To demonstrate improvements in sample efficiency we evaluate on multipl…
Vrije Universiteit Amsterdam（阿姆斯特丹自由大学）；University of Amsterdam（阿姆斯特丹大学）
Although group convolutional networks are able to learn powerful representations based on symmetry patterns, they lack explicit means to learn meaningful relationships among them (e.g., relative positions and poses). In this paper, we present attentive group equivariant convolutions, a generalization of the group convolution, in which attention is applied during the course of convolution to accentuate meaningful symmetry combinations and suppress non-plausible, misleading ones. We indicate that prior work on visual attention can be described as special cases of our proposed framework and show …
Equivariant-Aware Structured Pruning for Efficient Edge Deployment: A Comprehensive Framework with Adaptive Fine-Tuning
Faculty of Computers and Information Technology, University of Tabuk（计算机与信息科技学院，塔布克大学）；Artificial Intelligence and Sensing Technologies (AIST) Research Center, University of Tabuk（人工智能与传感技术（AIST）研究中心，塔布克大学）
Frame-based Equivariant Diffusion Models for 3D Molecular Generation
University of Amsterdam（阿姆斯特丹大学）；AMLab；AI4Science Lab（AI4Science实验室）；Korteweg-de Vries Institute for Mathematics University of Amsterdam（Korteweg-de Vries数学研究所阿姆斯特丹大学）
University of California San Diego（加州大学圣地亚哥分校）；Northeastern University（东北大学）
Exploiting symmetry in dynamical systems is a powerful way to improve the generalization of deep learning. The model learns to be invariant to transformation and hence is more robust to distribution shift. Data augmentation and equivariant networks are two major approaches to injecting symmetry into learning. However, their exact role in improving generalization is not well understood. In this work, we derive the generalization bounds for data augmentation and equivariant networks, characterizing their effect on learning in a unified framework. Unlike most prior theories for the i.i.d. setting…
Stanford University（斯坦福大学）；Florida State University（佛罗里达州立大学）
We introduce a theoretical framework for performing statistical tasks---including, but not limited to, averaging and principal component analysis---on the space of (possibly asymmetric) matrices with arbitrary entries and sizes. This is carried out under the lens of the Gromov-Wasserstein (GW) distance, and our methods translate the Riemannian framework of GW distances developed by Sturm into practical, implementable tools for network data analysis. Our methods are illustrated on datasets of letter graphs, asymmetric stochastic blockmodel networks, and planar shapes viewed as metric spaces. On…
TraceGrad: a Framework Learning Expressive SO(3)-equivariant Non-linear Representations for Electronic-Structure Hamiltonian Prediction
Institute of Artificial Intelligence, Hefei Comprehensive National Science Center（合肥综合性国家科学中心人工智能研究院）；University of Science and Technology of China（中国科学技术大学）
Recent work in equivariant deep learning bears strong similarities to physics. Fields over a base space are fundamental entities in both subjects, as are equivariant maps between these fields. In deep learning, however, these maps are usually defined by convolutions with a kernel, whereas they are partial differential operators (PDOs) in physics. Developing the theory of equivariant PDOs in the context of deep learning could bring these subjects even closer together and lead to a stronger flow of ideas. In this work, we derive a $G$-steerability constraint that completely characterizes when a …
We demonstrate that a state-of-the art multi-grid preconditioner can be learned efficiently by gauge-equivariant neural networks. We show that the models require minimal re-training on different gauge configurations of the same gauge ensemble and to a large extent remain efficient under modest modifications of ensemble parameters. We also demonstrate that important paradigms such as communication avoidance are straightforward to implement in this framework.
The big empirical success of group equivariant networks has led in recent years to the sprouting of a great variety of equivariant network architectures. A particular focus has thereby been on rotation and reflection equivariant CNNs for planar images. Here we give a general description of $E(2)$-equivariant convolutions in the framework of Steerable CNNs. The theory of Steerable CNNs thereby yields constraints on the convolution kernels which depend on group representations describing the transformation laws of feature spaces. We show that these constraints for arbitrary group representations…
