Steerable Partial Differential Operators for Equivariant Neural Networks

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 PDO between feature vector fields is equivariant, for arbitrary symmetry groups G. We then fully solve this constraint for several important groups. We use our solutions as equivariant drop-in replacements for convolutional layers and benchmark them in that role. Finally, we develop a framework for equivariant maps based on Schwartz distributions that unifies classical convolutions and differential operators and gives insight about the relation between the two. Figure 1: A vector field (left) can be mapped to a scalar field (right) by applying certain partial differential operators (PDOs), such as the Laplacian of the divergence and the 2D curl. Such a PDO from a 2D vector to a scalar field can be represented as a 2 × 1 matrix, where each of the two entries is a one-dimensional PDO that acts on one of the two components of the vector field. Similarly, matrices of PDOs with different dimensions map between other types of fields. Our goal is to find all PDOs for which this map becomes equivariant, for arbitrary types of fields. For the implementation, we will later discretize PDOs as stencils (middle).