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 authors propose to enforce soft equivariance constraints in ANN layers through a Gaussian prior on layer weights. Experimental results on a variety of problems support the benefit of the approach perhaps unsurprisingly. The overall contribution feels quite straightforward and incremental.