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 . 1 INTRODUCTION . In many machine learning tasks , the data exhibits certain symmetries , such as translation- and sometimes rotation-invariance in image classification . To exploit those symmetries , equivariant neural networks have been widely studied and successfully applied in the past years , beginning with Group convolutional neural networks ( Cohen & Welling , 2016 ; Weiler et al. , 2018b ) . A significant generalization of Group convolutional networks is given by steerable CNNs ( Cohen & Welling , 2017 ; Weiler et al. , 2018a ; Weiler & Cesa , 2019 ) , which unify many different pre-existing equivariant models ( Weiler & Cesa , 2019 ) . They do this by representing features as fields of feature vectors over a base space , such as R2 in the case of two-dimensional images . Layers are then linear equivariant maps between these fields . This is very reminiscent of physics . There , fields are used to model particles and their interactions , with physical space or spacetime as the base space . The maps between these fields are also equivariant , with the symmetries being part of fundamental physical laws . It is also noteworthy that these symmetries are largely ones that appear the most often in deep learning , such as translation and rotation equivariance . These similarities have already led to ideas from physics being applied in equivariant deep learning ( Lang & Weiler , 2021 ) . However , one remaining difference is that physics uses equivariant partial differential operators ( PDOs ) to define maps between fields , such as the gradient or Laplacian . Therefore , using PDOs instead of convolutions in deep learning would complete the analogy to physics and could lead to even more transfer of ideas between subjects . Equivariant PDO-based networks have already been designed in prior work ( Shen et al. , 2020 ; Smets et al. , 2020 ; Sharp et al. , 2020 ) . Most relevant for our work are PDO-eConvs ( Shen et al. , 2020 ) , which can be seen as the PDO-analogon of group convolutions . However , PDO-eConvs are only one instance of equivariant PDOs and do not cover the most common PDOs from physics , such as the gradient , divergence , etc . Very similarly to how steerable CNNs ( Cohen & Welling , 2017 ; Weiler et al. , 2018a ; Weiler & Cesa , 2019 ) generalize group convolutions , we generalize PDO-eConvs by characterizing the set of all translation equivariant PDOs between feature fields over Euclidean spaces . Because of this analogy , we dub these equivariant differential operators steerable PDOs . These steerable PDOs and their similarity to steerable CNNs also raise the question of how equivariant PDOs and kernels relate to each other , and whether they can be unified . We present a framework for equivariant maps that contains both steerable PDOs and convolutions with steerable kernels as special cases . We then prove that this framework defines the most general set of translation equivariant , linear , continuous maps between feature fields , complementing recent work ( Aronsson , 2021 ) that describes when equivariant maps are convolutions . Since formally developing this framework requires the theory of Schwartz distributions , we cover it mainly in Appendix E , and the main paper can be read without any knowledge of distributions . However , we reference the main results from this framework in the paper where appropriate . In order to make steerable PDOs practically applicable , we describe an approach to find complete bases for vector spaces of equivariant PDOs and then apply this method to the most important cases . We have also implemented steerable PDOs for all subgroups of O ( 2 ) ( see supplementary material for the code ) . Our code extends the E2CNN library1 ( Weiler & Cesa , 2019 ) , which will allow practitioners to easily use both steerable kernels and steerable PDOs within the same library , and even to combine both inside the same network . Finally , we test our approach empirically by comparing steerable PDOs to steerable CNNs . In particular , we benchmark different discretization methods for the numerical implementation . In summary , our main contributions are as follows : • We develop the theory of equivariant PDOs on Euclidean spaces , giving a practical characterization of precisely when a PDO is equivariant under any given symmetry . • We unify equivariant PDOs and kernels into one framework that provably contains all translation equivariant , linear , continuous maps between feature spaces . • We describe a method for finding bases of the vector spaces of equivariant PDOs , and provide explicit bases for many important cases . • We benchmark steerable PDOs using different discretization procedures and provide an implementation of steerable PDOs as an extension of the E2CNN library . 1.1 RELATED WORK . Equivariant convolutional networks Our approach to equivariance follows the one taken by steerable CNNs ( Cohen & Welling , 2017 ; Weiler et al. , 2018a ; Weiler & Cesa , 2019 ; Brandstetter et al. , 2021 ) . They represent each feature as a map from the base space , such as Rd , to a fiber Rc that is equipped with a representation ρ of the point group G. Compared to vanilla CNNs , which have fiber R , steerable CNNs thus extend the codomain of feature maps . A different approach is taken by group convolutional networks ( Cohen & Welling , 2016 ; Hoogeboom et al. , 2018 ; Weiler et al. , 2018b ) . They represent each feature as a map from a group H acting on the input space to R. Because the input to the network usually does not lie in H , this requires a lifting map from the input space to H . Compared to vanilla CNNs , group convolutional networks can thus be understood as extending the domain of feature maps . When H = Rd oG is the semidirect product of the translation group and a pointwise group G , then group convolutions on H are equivalent to G-steerable convolutions with regular representations . For finite G , the group convolution over G simply becomes a finite sum . LieConvs ( Finzi et al. , 2020 ) describe a way of implementing group convolutions even for infinite groups by using a Monte Carlo approximation for the convolution integral . Steerable CNNs with regular representations would have to use similar approximations for infinite groups , but they can instead also use ( non-regular ) finite-dimensional representations . Both the group convolutional and the steerable approach can be applied to non-Euclidean input spaces—LieConvs define group convolutions on arbitrary Lie groups and steerable convolutions can be defined on Riemannian manifolds ( Cohen et al. , 2019b ; Weiler et al. , 2021 ) and homogeneous spaces ( Cohen et al. , 2019a ) . One practical advantage of the group convolutional approach employed by LieConvs is that it doesn ’ t require solving any equivariance constraints , which tends to make implementation of new groups easier . They also require somewhat less heavy theoretical machinery . On the other hand , steerable 1https : //quva-lab.github.io/e2cnn/ CNNs are much more general . This makes them interesting from a theoretical angle and also has more practical advantages ; for example , they can naturally represent the symmetries of vector field input or output . Since our focus is developing the theory of equivariant PDOs and the connection to physics , where vector fields are ubiquitous , we are taking the steerable perspective in this paper . Equivariant PDO-based networks The work most closely related to ours are PDO-eConvs ( Shen et al. , 2020 ) , which apply the group convolutional perspective to PDOs . Unlike LieConvs , they are not designed to work with infinite groups . The steerable PDOs we introduce generalizes PDO-eConvs , which are obtained as a special case by using regular representations . A different approach to equivariant PDO-based networks was taken by Smets et al . ( 2020 ) . Instead of applying a differential operator to input features , they use layers that map an initial condition for a PDE to its solution at a fixed later time . The PDE has a fixed form but several learnable parameters and constraints on these parameters—combined with the form of the PDE—guarantee equivariance . Sharp et al . ( 2020 ) also use a PDE , namely the diffusion equation , as part of their DiffusionNet model , which can learn on 3D surfaces . But instead of exploiting the connection between PDOs and convolutions , they solve the PDE in the spectral domain . Interestingly , the time evolution operator for the diffusion equation is exp ( t∆ ) , which can be interpreted as an infinite power series in the Laplacian , very reminiscent of the finite Laplacian polynomials that naturally appear throughout this paper . Studying the equivariance of such infinite series of PDOs might be an interesting direction for future work . 2 STEERABLE PDOS In this section , we develop the theory of equivariant PDOs . We will represent all features as smooth fields f : Rd → Rc that associate a feature vector f ( x ) ∈ Rc , called the fiber at x , to each point x ∈ Rd . We write Fi = C∞ ( Rd , Rci ) for the space of these fields f in layer i. Additionally , we have a group of transformations acting on the input space Rd , which describes under which symmetries we want the PDOs to be equivariant . We will always use a group of the form H = ( Rd , + ) oG , for some G ≤ GL ( d , R ) . Here , ( Rd , + ) refers to the group of translations of Rd , while G is some group of linear invertible transformations . The full group of symmetries H is the semidi- rect product of these two , meaning that each element h ∈ H can be uniquely written as h = tg , where t ∈ Rd is a translation and g ∈ G a linear transformation . For example , if G = { e } is the trivial group , we consider only equivariance under translations , as in classical CNNs , while for G = SO ( d ) we additionally consider rotational equivariance . Each feature space Fi has an associated group representation ρi : G→ GL ( ci , R ) , which determines how each fiber Rci transforms under transformations of the input space . Briefly , ρi associates an invertible matrix ρi ( g ) to each group element g , such that ρi ( g ) ρi ( g′ ) = ρi ( gg′ ) ; more details on representation theory can be found in Appendix B . To see why these representations are necessary , consider the feature space F = C∞ ( R2 , R2 ) and the group G = SO ( 2 ) . The two channels could simply be two independent scalar fields , meaning that rotations of the input move each fiber but do not transform the fibers themselves . Formally , this would mean using trivial representations ρ ( g ) = 1 for both channels . On the other hand , the two channels could together form a vector field , which means that each fiber would need to itself be rotated in addition to being moved . This would correspond to the representation ρ ( g ) = g. These two cases are visualized in Fig . 1 . In general , the transformation of a feature f ∈ Fi under an input transformation tg with t ∈ Rd and g ∈ G is given by ( ( tg ) Bi f ) ( x ) : = ρi ( g ) f ( g −1 ( x− t ) ) . ( 1 ) The g−1 ( x− t ) term moves each fiber spatially , whereas the ρi ( g ) is responsible for the individual transformation of each fiber . For a network , we will need maps between adjacent feature spaces Fi and Fi+1 . Since during this theory section , we only consider single layers in isolation , we will drop the index i and simply denote the layer map as Φ : Fin → Fout . We are particularly interested in equivariant maps Φ , i.e . maps that commute with the action of H on the feature spaces : Φ ( hBin f ) = hBout Φ ( f ) ∀h ∈ H , f ∈ Fin . ( 2 ) We call Φ translation-equivariant if Eq . ( 2 ) holds for h ∈ ( Rd , + ) , i.e . for pure translations . Analogously , Φ is G-equivariant if it holds for linear transformations h ∈ G. Because H is the semidirect product of Rd and G , a map is H-equivariant if and only if it is both translation- and G-equivariant .

Summary: This paper introduces equivariant Partial Differential Operators as a drop in replacement for equivariant kernels in $G$-steerable networks. The paper provides a thorough treatment on the subject matter and gives the necessary equivariant constraints for PDO's in a similar flavor to $G$-steerable kernels. In addition, explicit solutions to these constraints are provided for subgroups of $O(2)$, $O(3)$ and $SO(3)$. Finally, the paper also empirically validates equivariant PDO's on rotation MNIST and STL-10 highlighting important discussion points on the discretization error.