Residual Pathway Priors for Soft Equivariance Constraints

1 Introduction . Central to the expanding application of deep learning to structured data like images , text , audio , sets , graphs , point clouds , and dynamical systems , has been a search for finding the appropriate set of inductive biases to match the model to the data . These inductive biases , such as recurrence [ 43 ] , local connectivity [ 29 ] , equivariance [ 10 ] , or differential equations [ 8 ] , reduce the set of explored hypotheses and improve generalization . Equivariance in particular has had a large impact as it allows ruling out a large class of meaningless shortcut features in many distinct domains , such as the ordering of the nodes in graphs and sets or the coordinate system chosen for an image . A disadvantage of hard coding these restrictions is that this prior knowledge may not match reality . A scene may have long range non-local interactions , rotation equivariance may be violated by a preferred camera angle , or a dynamical system may occasionally have discontinuous transitions . In particular , symmetries are delicate . A small perturbation like adding wind breaks the rotational symmetry of a pendulum , and bumpy or tilted terrain could break the translation symmetry for locomotion . In these cases we would like to incorporate our prior knowledge in a way that admits our own ignorance , and allows for the possibility that the world is more complex than we imagined . We aim to develop an approach that is more general , and can be applied when symmetries are exact , approximate , or non-existent . The Bayesian framework provides a mechanism for expressing such knowledge through priors . In much of the past work on Bayesian neural networks , the relationship between the prior distribution and the functions preferred by the prior is not transparent . While it is easy to specify different variances for different channels , or to use heavy tailed distributions , it is hard know how high level properties meaningfully translate into these low level attributes . Ultimately priors should represent our prior beliefs , and those beliefs we have are about high level concepts like the locality , independence , and symmetries of the data . ∗Equal Contribution 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . To address the need for more interpretable priors we introduce Residual Pathway Priors ( RPPs ) , a method for converting hard architectural constraints into soft priors . Practically , RPPs allow us to tackle problems in which perfect symmetry has been violated , but approximate symmetry is still present , as is the case for most real world physical systems . RPPs have a prior bias towards equivariant solutions , but are not constrained to equivariance . The choice of RPPs can be viewed from an Occam ’ s razor perspective , shown in Figure 1a , in which we seek to use models that have both the correct inductive biases and level of flexibility [ 41 ] . As we find with problems in which symmetries exist , highly flexible models with weak inductive biases like MLPs fail to concentrate prior mass around solutions that exhibit any symmetry . On the other hand when symmetries are only approximate , the strong biases of constrained models like Equivariant Multi-Layer Perceptrons ( EMLP ) [ 14 ] fail to provide support for the observations . As a middle ground between these two extremes , RPPs combine the inductive biases of constrained models with the flexibility of MLPs to define a model class which excels when data show approximate symmetries . In the following sections we introduce our method and show results across a variety of domains . We list our contributions and the accompanying sections below : 1 . We propose Residual Pathway Priors as a mechanism to imbue models with soft inductive biases , without constraining flexibility . 2 . While our approach is general , we use RPPs to show how to turn hard architectural constraints into soft equivariance priors ( Section 4 ) . 3 . We demonstrate that RPPs are robust to varying degrees of symmetry ( Section 5 ) . RPPs perform well under exact , approximate , or misspecified symmetries . 4 . Using RPP on the approximate symmetries in the complex state spaces of the Mujoco locomotion tasks , we improve the performance of model free RL agents ( Section 6 ) . We provide a PyTorch implementation of residual pathway priors at https : //github.com/mfinzi/residual-pathway-priors . 2 Related Work . The challenge of equivariant models not being able to fully fit the data has been identified in a number of different contexts , and with different application specific adjustments to mitigate the problem . Liu et al . [ 32 ] observe that convolutional networks can be extremely poor at tasks that require identifying or outputting spatial locations in an image as a result of the translation symmetry . The authors solve the problem by concatenating a coordinate grid to the input of the convolution layer . Constructing translation and rotation equivariant GCNNs , Weiler and Cesa [ 53 ] find that in order to get the best performance on CIFAR-10 and STL-10 datasets which have a preferred camera orientation , they must break the symmetry which they do by using equivariance to progressively smaller subgroups in the later layers . Bogatskiy et al . [ 6 ] go to great lengths to construct Lorentz group equivariant networks for tagging collisions in particle colliders only to break it by introducing dummy inputs that identify the collision axis . van der Wilk et al . [ 50 ] use the marginal likelihood to learn approximate invariances in Gaussian processes from data . In a similar approach , Benton et al . [ 5 ] explore a related concept of learning the extent of the symmetry using the reparametrization trick and test time augmentation . While sharing some commonalities with RPP , this method is not aimed at achieving approximate equivariance and can not bake equivariance into the model architecture . Outside of equivariance , adding model outputs to a much more restrictive base model has been a fruitful idea employed in multiple contexts . The original ResNet [ 20 , 21 ] drew on this motivation , with shortcut connections . Johannink et al . [ 24 ] , Silver et al . [ 45 ] proposed Residual Reinforcement Learning , whereby the RL problem is split into a user designed controller using engineering principles and a flexible neural network policy learned with RL . Similarly , in modeling dynamical systems , one approach is to incorporate a base parametric form informed by models from physics or biology , and only learn a neural network to fit the delta between the simple model and reality [ 25 , 33 ] . There have been several works tackling symmetries and equivariance in RL , such as permutation equivariance for multi agent RL [ 46 , 23 , 31 ] , as well exploring reflection symmetry for continuous control tasks [ 1 ] , and discrete symmetries in the more general framework of MDP homomorphisms [ 49 ] . However , in each of these applications the symmetries need to be exact , and the complexities of real data often require violating those symmetries . Although not constructed with this purpose , some methods which use regularizers to enforce equivariance [ 48 ] could be used for approximate symmetries . Interestingly , the value of approximate symmetries of MDPs has been explored in some theoretical work [ 42 , 47 ] , but without architectures that can make use of it . Additionally data augmentation , while not able to bake in architectural equivariance , has had successful application for encouraging equivariance on image tasks [ 27 ] and recently even on tabular state vectors [ 30 , 36 ] . 3 Background . In order develop our method , we first review the concept of group symmetries , how representations formalize the way these symmetries act on different objects . Group Symmetries In the machine learning context , a symmetry group G can be understood as a set of invertible transformations under which an object is the same , such as reflections or rotations . These symmetries can act on many different kinds of objects . A rotation could act on a simple vector , a 2d array like an image , a complex collection objects like the state space of a robot , or more abstractly on an entire classification problem or Markov Decision Process ( MDP ) . Representations The way that symmetries act on objects is described by a representation . Given an object in an n-dimensional vector space V , a group representation is a mapping ρ : G→ Rn×n , yielding a matrix which acts on V . Vectors v ∈ V are transformed v 7→ ρ ( g ) v. In deep learning , each of the inputs and outputs to our models can be embedded in some vector space : an m×m sized rgb image exists in R3m2 , and a node valued function on a graph of m elements exists within Rm . The representation ρ specifies how each of these objects transform under the symmetry group G. These representations can be composed of multiple simpler subrepresentations , describing how each object within collection transform . For example given the representation ρ1 of rotations acting on a vector in R3 and a representation ρ2 of how rotations act on a 3 × 3 matrix , the two objects concatenated together have a representation given by ρ1 ( g ) ⊕ ρ2 ( g ) = [ ρ1 ( g ) 0 0 ρ2 ( g ) ] where the two matrices are concatenated along the diagonal . Practically this means we can represent intricate and multifaceted structures by breaking them down into their component parts and defining how each part transforms . For example , we may know that the velocity vector , an orientation quaternion , a joint angle , and a control torque all transform in different ways under a left-right reflection , and one can accommodate this information into the representation . Equivariance Given some data X with representation ρin and Y with representation ρout , we may wish to learn some mapping f : X → Y . A model f is equivariant [ 10 ] , if applying the symmetry transformation to the input is equivalent to applying it to the output f ( ρin ( g ) x ) = ρout ( g ) f ( x ) . In other words , it is not the symmetry of X or Y that is relevant , but the symmetry of the function f mapping from X to Y . If the true relationship in the data has a symmetry , then constraining the hypothesis space to functions f that also have the symmetry makes learning easier and improves generalization [ 12 ] . Equivariant models have been developed for a wide variety of symmetries and data types like images [ 10 , 54 , 56 , 53 ] , sets [ 55 , 35 ] , graphs [ 34 ] , point clouds [ 3 , 15 , 44 ] , dynamical systems [ 13 ] , jets [ 6 ] , and other objects [ 51 , 14 ] . 4 Residual Pathway Priors . In this section , we introduce Residual Pathway Priors . The core implementation of RPP is to expand each layer in model into a sum of both a restrictive layer that encodes the hard architectural constraints and a generic more flexible layer but penalize the more flexible path via a lower prior probability . Through the difference in prior probability , explanations of the data using only the constrained solutions are prioritized by the model ; however , if the data is more complex the residual between the target and the constrained layer will be explained using the flexible layer . We can apply this procedure to any restriction priors , such as linearity , locality , markovian structure , and of course equivariance . The Residual Pathway Prior draws inspiration from the residual connections in ResNets [ 20 , 21 ] , whereby training stability and generalization improves by providing multiple paths for gradients to flow through the network that have different properties . One way of interpreting a residual block and shortcut connection f ( x ) = x + h ( x ) in combination with l2 regularization , either explicitly from weight decay or implicitly from the training dynamics [ 39 ] , is as a prior that places higher prior likelihood on the much simpler identity mapping than on the more flexible function h ( x ) . In this way , h ( x ) need only explain the the difference between what is explained in the previous layer ( passed through by I ) and the target . We can leverage a similar approach to convert the hard constraints of a restriction prior specified through a given network architecture into a soft prior that merely places higher prior likelihood on such models . Supposing we have an equivariant layer A ( x ) , and a more flexible non-equivariant layer B ( x ) which contains A as a special case , we can allow the model to explain part of the data with A and part with B by forming the sum A ( x ) + B ( x ) . A given a prior likelihood on the size of A and B such as p ( A ) ∝ exp ( −‖A‖2/2σ2a ) and p ( B ) ∝ exp ( −‖B‖2/2σ2b ) with σa > σb , a MAP optimized model will favor explanations of the data using the more structured layer A and only resort to using layer B to explain the difference between the target and what is already explained by the more structured model A . Adding these non-equivariant residual pathways to each layer of an equivariant model , we have a model that has the same expressivity of a network formed entirely of B layers , but with the inductive bias towards a model formed entirely with the equivariant A layers . We term this a residual pathway prior ( RPP ) . To make the approach concrete , we first consider constructing equivariance priors using the constraint solving approach known as Equivariant Multi-Layer Perceptrons ( EMLP ) from Finzi et al . [ 14 ] . Equivariant MLPs EMLPs provide a method for automatically constructing exactly equivariant layers for any given group and representation by solving a set of constraints . The way in which the vectors are equivariant is given by a formal specification of the types of the input and output through defining their representations . Given some input vector space Vin with representation ρin and some output space Vout with representation ρout the space of all equivariant linear layers mapping Vin → Vout satisfies ∀g ∈ G : ρout ( g ) W =Wρin ( g ) . These solutions to the constraint form a subspace of matrices Rnout×nin which can be solved for and described by a r dimensional orthonormal basis Q ∈ Rnoutnin×r . Linear layers can then be parametrized in this equivariant basis . The elements of W can be parametrized vec ( W ) = Qβ for β ∈ Rr for the linear layer v 7→Wv , and symmetric biases can be parametrized similarly . Equivariance Priors with EMLP In order to convert the hard equivariance constraints in EMLP into a soft prior over equivariance that can accomodate approximate symmetries , we can apply the RPP procedure from above to each these linear layers in the network . Instead of parametrizing the weights W directly in the equivariant basis vec ( W ) = Qβ , we can instead define W as the sum W = A + B of an equivariant weight matrix vec ( A ) = Qβ with a Gaussian prior over β and an unconstrained weight matrix B ∼ N ( 0 , σ2b I ) with a Gaussian prior . A Gaussian prior over β ∼ N ( 0 , σ2aI ) is equivalent to the Gaussian prior A ∼ N ( 0 , σ2aQQ > ) . Since Q is an orthogonal matrix , we can breakdown the covariance of B into the covariance in the equivariant subspace Q as well as the covariance in the orthogonal complement P , σ2b I = σ 2 bQQ > + σ2bPP > . Therefore the sum of the weight matrices is distributed as A+B =W ∼ N ( 0 , ( σ2a + σ2b ) QQ > + σ2bPP > ) . Regardless of the values of the prior variances σ2a and σ 2 b , solutions in the equivariant subspace QQ > are automatically favored by the model and assigned higher prior probability mass than those in the subspace PP > that violate the constraint . Even if σb > σa , the model still favors equivariance because the equivariance solutions are contained in the more flexible layer A . We show in Section 5.2 that RPPs are insensitive to the choice of σa and σb provided that σa is large enough to be able to fit the data . By replacing each of the equivariant linear layers in an EMLP with a sum of an equivariant layer and an unconstrained layer and adding in the negative prior likelihood to the loss function , we arrive at RPP-EMLP that can accommodate approximate or incorrectly specified symmetries . 1 RPPs With Other Equivariant Models While in EMLP these equivariant bases are solved for explicitly , the argument holds for the linear layers in other equivariant networks in precisely the same way . A good example of this is the translationally equivariant convolutional neural network ( CNN ) which can be viewed as a restricted subset of a fully connected network . Though the layers are parametrized as convolutions , the convolution operation can be expressed as a Toeplitz matrix residing within the space of dense matrices . Adding the convolution to a fully connected layer and choosing a prior variance σ2a and σ 2 b over each , we have the same RPP prior W ∼ N ( 0 , σ2aQQ > + σ2b I ) ( 1 ) where Q is the basis of ( bi- ) Toeplitz matrices corresponding to 3 × 3 filters . This RPP CNN has the biases of convolution but can readily fit non translationally equivariant data . Likewise for other equivariant models like GCNNs [ 10 ] , and we can even apply the RPP principle to the breaking of a given symmetry group to a subgroup .

Summary: The idea here is to use a combination of fully connected layers with equivariant layers and use a prior to bias the feedforward layer towards simpler (equivariant) linear maps. Standard residual connections are given as an extreme example where the simpler linear map is the identity operator. L2 penalty on the weights is used to implement this prior. Experiments study this type of model in settings where 1) the exact symmetry holds, 2)  it is partially violated, and 3) the data has no such symmetry. Example application domains used for experiments include physical models and applications in RL, where the correct (approximate) symmetries are known.