Title: Learning Layer-wise Equivariances Automatically using Gradients
Venue: NeurIPS 2023
Accepted: yes
Keywords: learning layer-wise relaxed equivariances bayesian symmetry discovery marginal likelihood
Abstract: Convolutions encode equivariance symmetries into neural networks leading to better generalisation performance. However, symmetries provide fixed hard constraints on the functions a network can represent, need to be specified in advance, and can not be adapted. Our goal is to allow flexible symmetry constraints that can automatically be learned from data using gradients. Learning symmetry and associated weight connectivity structures from scratch is difficult for two reasons. First, it requires efficient and flexible parameterisations of layer-wise equivariances. Secondly, symmetries act as constraints and are therefore not encouraged by training losses measuring data fit. To overcome these challenges, we improve parameterisations of soft equivariance and learn the amount of equivariance in layers by optimising the marginal likelihood, estimated using differentiable Laplace approximations. The objective balances data fit and model complexity enabling layer-wise symmetry discovery in deep networks. We demonstrate the ability to automatically learn layer-wise equivariances on image classification tasks, achieving equivalent or improved performance over baselines with hard-coded symmetry.

Learning Layer-wise Equivariances Automatically using Gradients

Abstract

Convolutions encode equivariance symmetries into neural networks leading to better generalisation performance. However, symmetries provide fixed hard constraints on the functions a network can represent, need to be specified in advance, and can not be adapted. Our goal is to allow flexible symmetry constraints that can automatically be learned from data using gradients. Learning symmetry and associated weight connectivity structures from scratch is difficult for two reasons. First, it requires efficient and flexible parameterisations of layer-wise equivariances. Secondly, symmetries act as constraints and are therefore not encouraged by training losses measuring data fit. To overcome these challenges, we improve parameterisations of soft equivariance and learn the amount of equivariance in layers by optimising the marginal likelihood, estimated using differentiable Laplace approximations. The objective balances data fit and model complexity enabling layer-wise symmetry discovery in deep networks. We demonstrate the ability to automatically learn layer-wise equivariances on image classification tasks, achieving equivalent or improved performance over baselines with hard-coded symmetry.

1 Introduction

Symmetry constraints, such as layer-wise equivariances of convolutional layers, allow neural networks to generalise efficiently [Cohen and Welling, 2016]. However, it is not always clear which symmetries to use and how much they should be enforced. In digit recognition, some rotational robustness is desirable, but strict invariance may make 6's and 9's indistinguishable. Similarly, translation equivariance prevents rate coding of relevant positional information [Sabour et al., 2017]. We propose a method to automatically learn the type and amount of layer-wise equivariances from training data.

To do so, we relax symmetries building upon recent literature on approximate equivariance, including residual pathways [Finzi et al., 2021a] and non-stationary filters [van der Ouderaa et al., 2022], effectively allowing controllable interpolation between strict equivariance or less constrained functions. To remain practical, we propose extensions to keep the parameter count similar to classic convolutions.

Allowing for differentiable selection between equivariances alone is not enough, as symmetries act as constraints and are therefore not encouraged by training losses that measure data fit, as noted in van der Ouderaa and van der Wilk [2021], Immer et al. [2021]. Search over many continuous hyperparameters with cross-validation is infeasible, relying on expensive retraining, hold-out validation data and does not leverage gradients. Instead, we follow Bayesian model selection through differentiable Laplace approximations to learn layer-wise equivariances from training data in a single training procedure.

We demonstrate automatically learning layer-wise symmetry structure on image classification tasks. To do so, we improve upon existing parameterisations of differentiable equivariance and derive corresponding Kronecker-factored Laplace approximations to the marginal likelihood. On image classification, we show that our method automatically learns convolutional structure in early layers and achieves similar or improved performance compared to architectures with hard-coded symmetry.

2 Related Work

Enforcing strict equivariance into the architecture. It can be shown that enforcing equivariance constraints in linear neural network layers is equivalent to performing a convolution [Kondor and Trivedi, 2018, Cohen et al., 2019]. Regular convolutional layers embed equivariance to translation, but this notion can be extended to other groups by convolving over groups [Cohen and Welling, 2016]. Group convolutions have been proposed for various domains and more complex continuous groups [Weiler et al., 2018, Weiler and Cesa, 2019, Worrall et al., 2017]. Equivariances enforce hard constraints on functions and can even form a linear subspace that can be explicitly computed [Van der Pol et al., 2020, Finzi et al., 2021b].

Place coding, rate coding, disentanglement and part-whole hierarchies. Equivariance closely relates to disentangled representations between place-coded and rate-coded features found in capsule networks [Sabour et al., 2017, Hinton et al., 2018, Kosiorek et al., 2019], as formalised in Cohen and Welling [2014], Cohen et al. [2018]. Positional information under strict translation equivariance of convolutions can only be place-coded, not rate-coded in features. Capsule networks aim to break this strict disentanglement to allow the embedding of part-whole hierarchies [Sabour et al., 2017]. We argue that our method solves the same problem by relaxing the disentanglement of strict equivariance and provide an automatic procedure to learn the extent to which information is exchanged and moved from place-coded to rate-coded features.

Approximate equivariance. Symmetries are often assumed known, must be fixed, and can not be adapted. This can be desirable if symmetries are known [Veeling et al., 2018], but becomes overly restrictive in other real-world tasks where symmetries or the appropriate amount of symmetry is not clear. Enforcing symmetry regardless can hamper performance [Liu et al., 2018b, Wang et al., 2022].

Approximate symmetry can prevent symmetry misspecification without losing the benefits that inductive biases of symmetries provide. Although strict group symmetries always apply to the entire group, by the group theoretical axiom of closure, approximate symmetry can be thought of as a form of robustness towards group actions that is only locally enforced (around training data). We use relaxations of symmetry that follow notions of approximate symmetry in literature [van der Wilk et al., 2018, Wang et al., 2020]. Although invariances can often easily be relaxed by simply averaging data augmentations [Benton et al., 2020, Schwöbel et al., 2021, van der Ouderaa and van der Wilk, 2021], parameterising approximate equivariance [Wang et al., 2022] can be more difficult. We build upon recent attempts, including partial equivariance [Romero and Lohit, 2021], residual pathways [Finzi et al., 2021a], and non-stationary filters [van der Ouderaa et al., 2022, van der Ouderaa and van der Wilk, 2023]. Existing relaxations of equivariance often introduce many additional parameters due to added linear layers [Finzi et al., 2021a] or use of hyper networks [Romero and Lohit, 2021, van der Ouderaa et al., 2022], making adoption in practice difficult. We strive to keep parameter counts close to classical convolutional layers, sometimes requiring further simplifying assumptions.

Symmetry discovery and objectives To learn symmetry, we relax layer-wise symmetries and thereby allow them to be differentiably learned with gradients, with strict symmetry as limiting case. Since adding symmetries does not directly improve training losses that rely on training fit, there is no direct encouragement of the objective function to use symmetries. To overcome this, some works have considered using validation data [Maile et al., 2022, Zhou et al., 2020]. Alternatively, Benton et al. [2020], Finzi et al. [2021a], Yang et al. [2023], van der Ouderaa et al. [2022] consider learning symmetry from training data only, but to do so require explicit regularisation that needs tuning. Some problems with this approach have been noted in Immer et al. [2022], which, like van der Wilk et al. [2018], proposes to use Bayesian model selection to learn invariances in a single training procedure, use differentiable Laplace approximations to remain tractable in deep neural networks. The method only considers invariances, which can be less efficient because it relies on learning parameters in the sampler instead of learning regularized symmetry in intermediary representations. We extend this line of work to enable automatic learning of layer-wise equivariances from training data.

Neural architecture search As layer-wise equivariances imply a convolutional architectural structure [Cohen et al., 2019], symmetry discovery can be viewed as a form of neural architecture search (NAS). Our approach uses relaxed symmetries to allow differentiable search over architectural structures, similar to DARTS [Liu et al., 2018a]. Instead of using validation data, however, our method uses Bayesian model selection to learn architecture from training data only.

We distinguish ourselves from large parts of AutoML literature in that we do not require expensive outer loops [Zoph et al., 2018], external agents [Cubuk et al., 2018], or even validation data [Lorraine et al., 2020, Yeh et al., 2022]. Our method can be interpreted as a one-shot architecture search using Laplace approximations, similar to [Zhou et al., 2019]. We demonstrate equivariance learning, but see no theoretical reason why the proposed objective should be limited to symmetry discovery, making it potentially broadly applicable in architecture search. For instance, in approaches that aim to learn network depth [Antorán et al., 2020], filter sizes [Romero and Zeghidour, 2023] and other layer types from training data. Although interesting future work, we focus on learning layer-wise equivariances.

Bayesian Model Selection We propose to treat layer-wise equivariances as hyperparameters and learn them with Bayesian model selection - a well-understood statistical procedure. Traditionally, model selection has only been available for models with tractable marginal likelihoods, such as Gaussian Processes [Rasmussen, 2003], but approximations for deep learning have also been developed, such as Variational Inference schemes [Ober and Aitchison, 2021], deep Gaussian Processes [Dutordoir et al., 2020], and scalable Laplace approximations [Immer et al., 2021], which we use. Although application in deep learning involves additional approximations, the underlying mechanism is similar to classic Automatic Relevance Determination [MacKay et al., 1994] to optimise exponential basis function lengthscales. Even more so, as we cast layer-wise equivariance learning as a lengthscale selection problem in Section 3.2. Learning invariances with Bayesian model selection in GPs was demonstrated in van der Wilk et al. [2018], and extensions to neural networks exist [Schwöbel et al., 2021, van der Ouderaa and van der Wilk, 2021], but often use lower bounds that only work on shallow networks. For deep networks, Immer et al. [2022, 2023] and Mlodozeniec et al. [2023] recently demonstrated invariance learning with marginal likelihood estimates in large ResNets.

This work describes a method that enables layer-wise learning of equivariance symmetries from training data. We formulate different priors consisting of hyperparameters that place a different mass on functions that obey equivariance symmetry constraints and functions that do not. The amount of symmetry becomes a hyperparameter we empirically learn from data using Bayesian model selection.

In Section 3, we discuss how to efficiently parameterise relaxed equivariance in convolutional architectures. We propose additional factorisations and sparsifications in Section 3.1 to bring the parameter count of relaxed equivariance closer to that of classical convolutions. In Section 3.2, we describe how the parameterisations can be used to specify the amount of equivariance in the prior.

In Section 4, we deal with the objective function and discuss approximate marginal likelihood for layer-wise relaxed equivariance. We rely on recent advancements in scalable linearised Laplace approximations that use KFAC approximations of the Hessian, and in Sections 5.1 and 5.2 derive KFAC for proposed relaxed equivariance layers to use them in conjunction with this objective.

3 Differentiable parameterisation of layer-wise equivariance symmetry

To learn symmetries from training data, we use the strategy of relaxing equivariance constraints through parameterisations allowing interpolation between general linear mappings and linear mappings that are constrained to be strictly equivariant. In our choice of relaxation, we strive to keep parameter counts as close as possible to classical convolutional layers to remain useful in practice. We describe parameterisations of linear, equivariant and relaxed equivariant layers for discrete 2dimensional translation symmetry Z 2 , as this is the most widely used symmetry in deep learning. All layers in this work can be naturally extended to other symmetry groups, as described in App. A.

Fully-connected layer A fully-connected layer forms a linear map between an input feature map x : [0 , C ] × Z 2 → R (spatially supported at intervals [0 , X ] × [0 , Y ] ), with C input channels outputting a C ′ channel feature map y : [0 , C ′ ] × Z 2 → R with the same spatial support:

where non-zero weights in w : Z 6 ↦→ R can be represented in C ′ CX 2 Y 2 parameters. This may be more clear if, instead of using real-valued functions to describe features and weights, we equivalently flatten x , y as vectors y = Wx and write w as a matrix W ∈ R C ′ XY × CXY . Yet, we stick to the form of Eq. (1) as this simplifies future factorisations and generalisations. Furthermore, the notation emphasises that other layers, such as convolutions, form special cases of a fully-connected layer.

Convolutional layer We might want to constrain a layer to be strict equivariant , so that translations in the input ( dx, dy ) ∈ Z 2 yield equivalent translations in the output, mathematically:

It can be shown that equivariance is equivalent (if and only if) to having a weight matrix that is stationary along the spatial dimension, meaning it solely depends on relative coordinates ¯ x = x ′ -x and ¯ y = y ′ -y and can thus be written in terms of stationary weights ¯ θ : ( c, c ′ , ¯ x, ¯ y ) ↦→ R :

with C ′ CXY parameters in ¯ θ . Note that Eq. (3) is exactly a convolutional layer ('convolution is all you need', Theorem 3.1 of Cohen et al. [2019]). Constraining the stationary weights ¯ θ to be only locally supported on a rectangular domain ¯ x, ¯ y ∈ [ -S -1 2 , S -1 2 ] , and 0 otherwise, gives rise to small (e.g. 3 × 3 ) convolutional filters and further reduces the required parameters to C ′ CS 2 .

3.1 Factoring and sparsifying relaxed equivariance

To learn layer-wise equivariances, we take the approach of relaxing symmetry constraints in neural network layers allowing explicit differentiable interpolation between equivariant and non-equivariant solutions. As a starting point, we consider the residual pathways of [Finzi et al., 2021a], reinterpreted as factorisation of the linear map as w = θ + ¯ θ with fully-connected θ : ( c ′ , c, x ′ , y ′ , x, y ) ↦→ R and stationary ¯ θ ( c ′ , c, ¯ x, ¯ y ) ↦→ R with ¯ x = x ′ -x and ¯ y = y ′ -y :

The layer consists of simply adding FC and CONV layer outputs. The CONV path spans the linear subspace of FC that obeys strict equivariance of Eq. (2). As such, the combined layer is equally expressive as a linear map (we can construct a bijection between w and θ + ¯ θ , as is common for relaxed equivariance [van der Ouderaa et al., 2022]). In Finzi et al. [2021a], independent priors are placed on both sets of weights with relative variances controlling how much equivariance is encouraged. We also consider this prior, but assume it is unknown and use Bayesian model selection to automatically learn the variances that control symmetry constraints from data. Further, naive residual pathways of Eq. (4) require polynomially many parameters C ′ CX 4 + C ′ CX 2 , or C ′ CX 4 + C ′ CS 2 for S × S filters, using square features X = Y for simplicity. This makes using residual pathways in common architectures infeasible (see line 3 in Table 2). We overcome this by proposing factorisations and sparsifications bringing parameter counts of relaxed equivariance closer to convolutional layers.

Factored layers Inspired by common factorisations in convolutional architectures, such as depthwise [Chollet, 2017] and group separability [Knigge et al., 2022], we propose to factor fully-connected layers between input and output dimensions θ ( c ′ , c, x ′ , y ′ , x, y ) = θ 1 ( c ′ , c, x ′ , y ′ ) θ 2 ( c ′ , c, x, y ) . This greatly reduces the parameters, while fully maintaining the benefits of fully-connected structure from full spatial H × W input and output H ′ × W ′ to both input and outputs channels C × C ′ .

The factorisation of Eq. (5) has several benefits. First, it reduces the number of parameters required for the expensive the fully-connected path from C ′ CX 2 Y 2 to 2 C ′ CXY . This is essential in practice to obtain parameter counts close to convolutional layers (133.6M → 1.5M in Table 2). Secondly, we can write the layer as a composition of two linear operations summing over ∑ x,y and ∑ c , which we use to derive Kronecker-factored approximations to the second-order curvature and enable the use of scalable Laplace approximations. Lastly, the parameters θ 1 and θ 2 have a single spatial domain that lends themselves to further optional sparsification with basis functions.

Spatial sparsification To further improve the parameter efficiency, we consider parameterising spatial dimensions in a lower dimensional basis ¯ s ω θ ( c ′ , c, ¯ x, ¯ y ) = ∑ M j =1 a j ϕ c ′ ,c j (¯ x, ¯ y ) . Our approach is partly inspired by B-spline Lie

group filters of [Bekkers, 2019]. We follow van der Ouderaa and van der Wilk [2023] proposing standard exponential basis functions ϕ c ′ ,c j (¯ x, ¯ y ) = u j exp( -ω 2 ([¯ x ¯ y ] T -z j ) 2 ) for convolutional filters, with anchor point locations z j ∈ R 2 and anchor point values u j ∈ R . The work demonstrated that lengthscales ω can allow differentiable control over equivariances (used in Section 3.2) as well as practical sparsification with negligible performance loss, even if filters are already small (e.g. 3 × 3 ).

This reduces the number of parameters from C ′ CS 2 to C ′ CP with P anchor points, or C ′ CP +2 P if anchor point locations { z j } j M =1 (shared over channels) are also considered learnable parameters. As weights in our proposed factored F-FC have independently factored spatial input and output domains, we can perform a similar sparsification. We write s ω 1 and s ω 2 for spatially sparsified θ 1 and θ 2 :

reducing the parameter cost from 2 C ′ CXY to C ′ CP with P total number of anchor points. Compared to convolutional filters, which already have small spatial supports, sparsification of fullyconnected layers will typically result in a much larger absolute reduction of parameters, for fixed P . As discussed in App. A, sparsified layers of Eqs. (6) and (7) can be extended to other groups, but require well-defined basis functions ϕ on groups, as explored in [Azangulov et al., 2022].

3.2 Specifying the amount of equivariance in the prior

To allow for learnable equivariance, we explicitly specify symmetry constraints in the prior and then empirically learn them through approximate Bayesian model selection. We choose a Gaussian form N ( 0 , σ 2 l I ) and treat variances as hyperparameters η = [ σ 1 , σ 2 , . . . ] . We carefully parameterise the network, such that the prior in each layer l enforces an equivariance symmetry in the following limit:

σ l = 0 = ⇒ strict equivariance (Eq. (2)) and σ l > 0 = ⇒ relaxed equivariance (8) We can think of different possible ways to parameterise neural network layers such that the conditions of Eq. (8) are met. For instance, Finzi et al. [2021a] propose placing priors directly on the weights of both residual pathways FC+CONV N ( θ | 0 , σ 2 l I ) and N ( ¯ θ | 0 , ¯ σ 2 l I ) with scalar variances σ 2 l , ¯ σ 2 l ∈ R and identity matrix I ∈ R | θ |×| θ | . Intuitively, all weights in non-equivariant paths become zero in the zero variance limit σ l =0 , resulting in strict equivariance. We consider the same prior for the factored parameterisation F-FC+CONV. Likewise, for sparsified layers, we can place the prior directly on anchor points N ( u | 0 , σ 2 l I ) and N (¯ u | 0 , ¯ σ 2 l I ) which will control equivariance through a similar mechanism. For S-CONV, this prior corresponds to the correlated prior of Fortuin et al. [2021] if anchor points are placed exactly on the filter grid, whereas our anchor points are sparse M<S 2 and have learnable positions z j . Alternatively, for sparsified S- layers, we can consider placing the prior directly on the lengthscales N ( ω | 0 , σ 2 l I ) and N (¯ ω l | 0 , ¯ σ 2 l I ) , as proposed in non-stationary filters [van der Ouderaa et al., 2022, van der Ouderaa and van der Wilk, 2023]. Intuitively, the limit of σ l =0 forces ω l to be zero, causing non-stationary components to be constant and weights to act stationary, resulting in strict equivariance. Treating lengthscales ω as learnable hyperparameter has the additional advantage that it allows learning of convolutional filter frequencies from data through model selection of hyperparameter ¯ σ 2 l .

Different constructions of the prior may induce different prior probabilities in function space. Nevertheless, it can be shown that Eq. (8) holds for all choices of prior described above. Thus, we have explicit control over symmetry with strict equivariance at the σ l =0 limiting case and relaxed layer-wise equivariance at higher variances σ l > 0 . Unlike prior works [Finzi et al., 2021a, van der Ouderaa et al., 2022] that require setting or tuning the prior variances σ 2 l controlling equivariance constraints, we follow Bayesian model selection to learn them automatically from training data.

4 Bayesian model selection of symmetry constraints

Inspired by Bayesian methodology, we propose to empirically learn the amount of layer-wise equivariance in a neural network from training data through approximate Bayesian model selection. We define equivariances in our prior and treat the amount of symmetry as hyperparameters optimised using the marginal likelihood - a well-understood statistical procedure known as 'empirical Bayes'. Mathematically, we find a point estimate of the hyperparameters η ∗ =argmax η p ( D| η ) , whilst integrating out the P model parameters θ ∈ R P :

The marginal likelihood p ( D| η ) forms the normalising constant of the unnormalised posterior over parameters:

As the marginal likelihood is constant in θ , it can be ignored when training neural network parameters with a regular MAP point-estimate θ ∗ = arg min θ L θ :

To optimise hyperparameters, the marginal likelihood p ( D| η ) can not be ignored and becomes important again as the objective function for hyperparameters. The marginal likelihood is the integral over the very high-dimensional space of possible parameter configurations, which is intractable for large neural networks and therefore requires additional approximations in practice.

To overcome the intractability of the marginal likelihood, the Laplace approximation [MacKay, 2003] can be used to estimate the normalising constant by first approximating the posterior by a Gaussian centered at the mode µ = θ ∗ and the covariance set to the local curvature Σ = H -1 θ ∗ := -∇ 2 θ L θ | θ = θ ∗ at that point:

Integrals of Gaussians are known in closed-form, and we can thus use the normalising constant of q to estimate the log marginal likelihood objective used to optimise hyperparameters η ∗ = arg min η L η :

Although the Laplace approximation greatly simplifies the marginal likelihood, the Hessian determinant | H θ ∗ | remains difficult to compute for large neural networks due to quadratic scaling in parameter count O ( | θ | 2 ) . In practice, we therefore use a local linearisation [Bottou et al., 2018] resulting in the generalized Gauss-Newton (GGN) approximation, which we further approximate using a block-diagonal Kronecker-factored (KFAC) structure Martens and Grosse [2015] (see Section 5).

The objective of Eq. (9) penalises model complexity enforcing a form of Occam's razor [Rasmussen and Ghahramani, 2000]. Complexity control is important when learning symmetry constraints, as previous works [van der Ouderaa and van der Wilk, 2021, Immer et al., 2021] have shown that just using the regular training loss L θ w.r.t. data augmentation parameters prefers collapse into a solution without symmetry constraints resulting in fully-connected structure. The marginal likelihood overcomes this issue and can learn symmetries from data, as shown in van der Wilk et al. [2018]. The Laplace approximated marginal likelihood L η can be viewed as a form of curvature-aware minimisation and has recently been used to learn invariances in large ResNets [Immer et al., 2022].

5 Kronecker-factored curvature of relaxed equivariance layers

In Eq. (9), we propose the Laplace approximated log marginal likelihood L η as objective to learn layer-wise equivariances from data. As Hessians scale quadratically in the number of parameters

O ( | θ | 2 ) , the loss becomes intractable for large networks. We use KFAC [Martens and Grosse, 2015] to approximate the Hessian of the log likelihood H θ ∗ and have the Gauss-Newton approximation H :

where the Jacobians of network outputs f with respect to parameters θ for data points x n are [ J ( x n )] k,p = ∂f k ∂θ p ( x n ) , and [ Λ ( x n )] k,g = ∂ log p ( y n | x n ; f ( x n ; θ )) ∂f k ∂f g is the log likelihood Hessian, all evaluated at θ ∗ . KFAC is a block-diagonal approximation of the GGN, with individual blocks H l = ∑ n H n,l = ∑ n J l ( x n ) ⊤ Λ ( x n ) J l ( x n ) for each layer l with parameters θ l and Jacobian J l ( x ) .

For a fully-connected layer, θx with θ ∈ R G l × D l and x ∈ R D l , we can write Jacobians of a single data point w.r.t. parameters as J l ( x n ) ⊤ = a l,n ⊗ g l,n , where a l,n ∈ R D l × 1 is the input to the layer and g l,n ∈ R G l × K is the transposed Jacobian of f w.r.t. the l th layer output. We thus have:

KFAC only requires the computation of the two last Kronecker factors, which are relatively small compared to the full GGN block. KFAC has been extended to convolutional CONV layers in Grosse and Martens [2016]. We derive KFAC also for factored layers F-FC of Eq. (5) and sparsified S-FC and S-CONV layers of Eqs. (6) and (7), allowing use in scalable Laplace approximations to learn layer-wise equivariances from data. The essential step to extend KFAC is to write the Jacobian of the l th layer, J l ( x ) , as Kronecker product allowing data points to be factored and approximated with the last KFAC step. Our derivations of KFAC for new layers are in detail provided in Apps. C.1 and C.2.

5.1 KFAC for factored layers

To obtain the KFAC of F-FC in Eq. (5), we write the application of both weight parameters θ 1 and θ 2 as a matrix mutiplication akin to a linear layer, which allows us to apply a similar approximation as in standard KFAC. In particular, we first split up the operations in Eq. (5) into first applying θ 2 and then θ 1 . Each of these operations can be written as a matrix product θx with corresponding appropriate sizes of θ and x . In comparison to the standard KFAC case, we have a Jacobian that can be written as

which means an additional sum needs to be handled. In particular, plugging the Jacobian structure into the standard KFAC approximation, we have

where the first approximation is from the Hessian to the GGN and the second approximation is due to KFACs exchange of sums and products in favour of numerical efficiency. The efficiency comes due to the fact that we have two Kronecker factors instead of a single dense matrix. This allows to compute and optimize the marginal likelihood approximation efficiently and therefore tune σ l values. In App. C.1, we derive the KFAC for θ 1 and θ 2 in detail.

5.2 KFAC for sparsified layers

To obtain KFAC for sparsified S-FC and S-CONV layers, we extend derivations of KFAC for respective factored F-FC layers of Section 5.1 and CONV layers in [Grosse and Martens, 2016]. To compute Jacobians with respect to our new anchor points a parameters, instead of the induced weights s ω , we apply the chain rule:

Consequently, we can compute KFAC for sparsified layers by projecting the appropriate dimensions of Kronecker factors of existing KFAC derivations with basis functions ϕ stored in the forward pass. In App. C.2, we derive KFAC for sparsified S-FC and S-CONV layers in detail.

6 Experiments

6.1 Toy problem: adapting symmetry to task

The symmetries in a neural network should be task dependent. Imagine a task where the full symmetry would be too restrictive. Can our method automatically adapt to what is required by the task? Although this ability arguably becomes most important on more complex datasets where symmetries are not strict or completely unknown, the absence of ground-truth symmetry in such scenarios makes evaluation more difficult. We, therefore, turn to a simple toy problem of which underlying symmetries are known and consider more interesting larger-scale datasets in the following sections. We slightly altered the MNIST dataset [LeCun et al., 1989] so that digits are randomly placed in one of four image quadrants (see example left of Table 1) and consider the task of classifying the correct digit out of 10 classes, and the task of classifying both the digit and the quadrant it is placed in, totalling 40 classes. The first task is translation invariant, as moving a digit around does not change its class. The second task, on the other hand, can not be solved under strict translation invariance due to the dependence on location. As expected, we find in Table 1 that the CONV model outperforms the FC model on the first strictly symmetric task. Conversely, the CONV model fails to solve the second task due to the symmetry misspecification and is consequently outperformed by the FC model. The proposed method with adjustable symmetry achieves high test accuracy on both tasks.

(3, upper left)

6.2 Learning to use layer-wise equivariant convolutions on CIFAR-10

Convolutional layers provide useful inductive bias for image classification tasks, such as CIFAR-10 [Krizhevsky et al., 2009]. We investigate whether our method can select this favourable convolutional structure when presented with this task. We follow Finzi et al. [2021a] and use the architecture of Neyshabur [2020] (see App. F). In Table 2, we compare test performance when using fully-connected FC, convolutional CONV, residual pathways FC+CONV, and the proposed more parameter-efficient factorisations F-FC and F-FC+CONV and sparsified S- layers. The prior is placed on the weights θ , anchor points u , or lengthscales ω , as discussed in Section 3.2. As expected, CONV layers perform best under regular MAP training, followed by less constrained FC and FC+CONV layers. If we, instead, use the Laplace approximated marginal likelihood estimates, we find that learnable symmetry obtains equivalent (within <1%) performance as the convolutional CONV models.

Layers learn to become convolutional, except at the end of the network.

Upon further inspection of learned symmetries, we observe very high prior precisions σ -2 l > 10 6 for non-equivariant paths in most layers, thus negligibly contributing to final outputs (see App. D.1). We, therefore, conclude that layers have learned to become convolutional from training data, which also explains why training learned equivariance with Laplace performs on par with CONV models in Table 2. We reach the same conclusions upon inspection of effective dimensions [MacKay, 1992] in App. D.2 and visualise learned symmetries based on the highest relative effective number of parameters in Fig. 3. If layers did not become convolutional, this always occurred at the end of the network. This is particularly interesting as it coincides with common practice in literature. Spatial dimensions in the used architecture are halved through pooling operations, reaching 1 with sufficient or global pooling, after which features become strictly invariant and there is no difference anymore between 1x1 CONV and FC layers. Convolutional architectures in literature often flatten features and apply FC layers before this point [LeCun et al., 1998], breaking strict symmetry and allowing some dependence in the output on absolute position. Our findings indicate that the method can discover such favourable

x

CONV

32 x 32

16 x 16

8 x 8

4 x 4

2 x 2

1 x 1

Equivariance

in first layers

No equivariance

at the end of model

CONV

CONV

CONV

CONV

CONV

CONV

FC

FINAL

FINAL

Final layers respect

symmetry, because

of 1x1 spatial dims

y

Figure

3:

Learned

CONV+FC model.

network designs automatically from data, outperforming regular MAP training. Furthermore, it shows that not always the same layer is selected, but different layer types are learned on a per-layer basis.

6.3 Selecting symmetry from multiple groups

App. A generalises the proposed method to other symmetry groups, and App. B describes layer-wise equivariance selection from a set of multiple symmetry groups or layer types. In this experiment, we consider adding discrete 90-degree rotations to the architecture F-FC+CONV+GCONV. In Table 3, we compare the test performance with MAP and approximated marginal likelihood training on versions of MNIST and CIFAR-10 datasets (details in App. F.5). We compensate for the added parameters resulting from the rotationally equivariant GCONV path by reducing channel sizes of individual paths by a factor of 5 ( α =10 to α =2 , see App. F). Still, we are able to obtain 80% test accuracy on CIFAR-10 when trained with approximated marginal likelihood. Upon analysis of learned prior variances (App. D.1) and effective number of parameters (App. D.2), we observe a positive correlation between learned symmetries and augmentations applied to the datasets trained on.

7 Discussion and Conclusion

This work proposes a method to automatically learn layer-wise equivariances in deep learning using gradients from training data. This is a challenging task as it requires both flexible parameterisations of symmetry structure and an objective that can learn symmetry constraints. We improve upon existing parameterisations of relaxed equivariance to remain practical in the number of parameters. We learn equivariances through Bayesian model selection, by specifying symmetries in the prior and learning them by optimising marginal likelihood estimates. We derive Kronecker-factored approximations of proposed layers to enable scalable Laplace approximations in deep neural networks.

The approach generalises symmetry groups and can be used to automatically determine the most relevant symmetry group from a set of multiple groups. We rely on relaxing symmetries and learning the amount of relaxation, where strict equivariance forms a limiting case. In doing so, the method does require a list of symmetry groups and associated group representations that are considered, and learns to ignore, partially use, or strictly obey symmetry constraints. Yet, we argue this is a huge step forward compared to always using a single, strict symmetry. We hope that allowing artificial neural networks to automatically adapt structure from training data helps to leverage existing geometry in data better and reduce the reliance on trial-and-error in architecture design.

References

Javier Antorán, James Allingham, and José Miguel Hernández-Lobato. Depth uncertainty in neural networks. Advances in neural information processing systems , 33:10620-10634, 2020.

Iskander Azangulov, Andrei Smolensky, Alexander Terenin, and Viacheslav Borovitskiy. Stationary kernels and gaussian processes on lie groups and their homogeneous spaces i: the compact case. arXiv preprint arXiv:2208.14960 , 2022.

Erik J Bekkers. B-spline cnns on lie groups. arXiv preprint arXiv:1909.12057 , 2019.

Gregory Benton, Marc Finzi, Pavel Izmailov, and Andrew Gordon Wilson. Learning invariances in neural networks. arXiv preprint arXiv:2010.11882 , 2020.

Léon Bottou, Frank E Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM review , 60(2):223-311, 2018.

Anadi Chaman and Ivan Dokmanic. Truly shift-invariant convolutional neural networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 3773-3783, 2021.

François Chollet. Xception: Deep learning with depthwise separable convolutions. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 1251-1258, 2017.

Taco Cohen and Max Welling. Learning the irreducible representations of commutative lie groups. In International Conference on Machine Learning , pages 1755-1763. PMLR, 2014.

Taco Cohen and Max Welling. Group equivariant convolutional networks. In International conference on machine learning , pages 2990-2999. PMLR, 2016.

Taco S Cohen, Mario Geiger, and Maurice Weiler. Intertwiners between induced representations (with applications to the theory of equivariant neural networks). arXiv preprint arXiv:1803.10743 , 2018.

Taco S Cohen, Mario Geiger, and Maurice Weiler. A general theory of equivariant cnns on homogeneous spaces. Advances in neural information processing systems , 32, 2019.

Ekin D Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc V Le. Autoaugment: Learning augmentation policies from data. arXiv preprint arXiv:1805.09501 , 2018.

Nima Dehmamy, Robin Walters, Yanchen Liu, Dashun Wang, and Rose Yu. Automatic symmetry discovery with lie algebra convolutional network. Advances in Neural Information Processing Systems , 34:2503-2515, 2021.

Vincent Dutordoir, Mark Wilk, Artem Artemev, and James Hensman. Bayesian image classification with deep convolutional gaussian processes. In International Conference on Artificial Intelligence and Statistics , pages 1529-1539. PMLR, 2020.

Marc Finzi, Gregory Benton, and Andrew G Wilson. Residual pathway priors for soft equivariance constraints. Advances in Neural Information Processing Systems , 34, 2021a.

Marc Finzi, Max Welling, and Andrew Gordon Wilson. A practical method for constructing equivariant multilayer perceptrons for arbitrary matrix groups. In International Conference on Machine Learning , pages 3318-3328. PMLR, 2021b.

Vincent Fortuin, Adrià Garriga-Alonso, Florian Wenzel, Gunnar Rätsch, Richard Turner, Mark van der Wilk, and Laurence Aitchison. Bayesian neural network priors revisited. arXiv preprint arXiv:2102.06571 , 2021.

Roger Grosse and James Martens. A kronecker-factored approximate fisher matrix for convolution layers. In International Conference on Machine Learning , pages 573-582. PMLR, 2016.

Geoffrey E Hinton, Sara Sabour, and Nicholas Frosst. Matrix capsules with em routing. In International conference on learning representations , 2018.

Alexander Immer, Matthias Bauer, Vincent Fortuin, Gunnar Rätsch, and Khan Mohammad Emtiyaz. Scalable marginal likelihood estimation for model selection in deep learning. In International Conference on Machine Learning , pages 4563-4573. PMLR, 2021.

Alexander Immer, Tycho F. A. van der Ouderaa, Vincent Fortuin, Gunnar Rätsch, and Mark van der Wilk. Invariance learning in deep neural networks with differentiable laplace approximations, 2022.

Alexander Immer, Tycho FA Van Der Ouderaa, Mark Van Der Wilk, Gunnar Ratsch, and Bernhard Schölkopf. Stochastic marginal likelihood gradients using neural tangent kernels. In International Conference on Machine Learning , pages 14333-14352. PMLR, 2023.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

David M Knigge, David W Romero, and Erik J Bekkers. Exploiting redundancy: Separable group convolutional networks on lie groups. In International Conference on Machine Learning , pages 11359-11386. PMLR, 2022.

Risi Kondor and Shubhendu Trivedi. On the generalization of equivariance and convolution in neural networks to the action of compact groups. In International Conference on Machine Learning , pages 2747-2755. PMLR, 2018.

Adam Kosiorek, Sara Sabour, Yee Whye Teh, and Geoffrey E Hinton. Stacked capsule autoencoders. Advances in neural information processing systems , 32, 2019.

Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.

Yann LeCun. The mnist database of handwritten digits. http://yann. lecun. com/exdb/mnist/ , 1998.

Yann LeCun, Bernhard Boser, John Denker, Donnie Henderson, Richard Howard, Wayne Hubbard, and Lawrence Jackel. Handwritten digit recognition with a back-propagation network. Advances in neural information processing systems , 2, 1989.

Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE , 86(11):2278-2324, 1998.

Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts: Differentiable architecture search. arXiv preprint arXiv:1806.09055 , 2018a.

Rosanne Liu, Joel Lehman, Piero Molino, Felipe Petroski Such, Eric Frank, Alex Sergeev, and Jason Yosinski. An intriguing failing of convolutional neural networks and the coordconv solution. arXiv preprint arXiv:1807.03247 , 2018b.

Jonathan Lorraine, Paul Vicol, and David Duvenaud. Optimizing millions of hyperparameters by implicit differentiation. In Silvia Chiappa and Roberto Calandra, editors, Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics , volume 108 of Proceedings of Machine Learning Research , pages 1540-1552. PMLR, 26-28 Aug 2020.

Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608.03983 , 2016.

David JC MacKay. Bayesian interpolation. Neural computation , 4(3):415-447, 1992.

David JC MacKay. Information theory, inference and learning algorithms . Cambridge university press, 2003.

David JC MacKay et al. Bayesian nonlinear modeling for the prediction competition. ASHRAE transactions , 100(2):1053-1062, 1994.

Kaitlin Maile, Dennis George Wilson, and Patrick Forré. Equivariance-aware architectural optimization of neural networks. In The Eleventh International Conference on Learning Representations , 2022.

James Martens. New insights and perspectives on the natural gradient method. The Journal of Machine Learning Research , 21(1):5776-5851, 2020.

James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curvature. In International conference on machine learning , pages 2408-2417, 2015.

Bruno Mlodozeniec, Matthias Reisser, and Christos Louizos. Hyperparameter optimization through neural network partitioning. arXiv preprint arXiv:2304.14766 , 2023.

Behnam Neyshabur. Towards learning convolutions from scratch. Advances in Neural Information Processing Systems , 33:8078-8088, 2020.

Sebastian W Ober and Laurence Aitchison. Global inducing point variational posteriors for bayesian neural networks and deep gaussian processes. In International Conference on Machine Learning , pages 8248-8259. PMLR, 2021.

Carl Rasmussen and Zoubin Ghahramani. Occam's razor. Advances in neural information processing systems , 13, 2000.

Carl Edward Rasmussen. Gaussian processes in machine learning. In Summer school on machine learning , pages 63-71. Springer, 2003.

David W Romero and Suhas Lohit. Learning equivariances and partial equivariances from data. arXiv preprint arXiv:2110.10211 , 2021.

David W Romero and Neil Zeghidour. Dnarch: Learning convolutional neural architectures by backpropagation. arXiv preprint arXiv:2302.05400 , 2023.

Sara Sabour, Nicholas Frosst, and Geoffrey E Hinton. Dynamic routing between capsules. Advances in neural information processing systems , 30, 2017.

Pola Schwöbel, Martin Jørgensen, Sebastian W Ober, and Mark van der Wilk. Last layer marginal likelihood for invariance learning. arXiv preprint arXiv:2106.07512 , 2021.

Tycho FA van der Ouderaa and Mark van der Wilk. Learning invariant weights in neural networks. In Workshop in Uncertainty & Robustness in Deep Learning, ICML , 2021.

Tycho FA van der Ouderaa and Mark van der Wilk. Sparse convolutions on lie groups. In NeurIPS Workshop on Symmetry and Geometry in Neural Representations , pages 48-62. PMLR, 2023.

Tycho FA van der Ouderaa, David W Romero, and Mark van der Wilk. Relaxing equivariance constraints with non-stationary continuous filters. arXiv preprint arXiv:2204.07178 , 2022.

Elise Van der Pol, Daniel Worrall, Herke van Hoof, Frans Oliehoek, and Max Welling. Mdp homomorphic networks: Group symmetries in reinforcement learning. Advances in Neural Information Processing Systems , 33:4199-4210, 2020.

Mark van der Wilk, Matthias Bauer, ST John, and James Hensman. Learning invariances using the marginal likelihood. arXiv preprint arXiv:1808.05563 , 2018.

Bastiaan S Veeling, Jasper Linmans, Jim Winkens, Taco Cohen, and Max Welling. Rotation equivariant cnns for digital pathology. In Medical Image Computing and Computer Assisted Intervention-MICCAI 2018: 21st International Conference, Granada, Spain, September 16-20, 2018, Proceedings, Part II 11 , pages 210-218. Springer, 2018.

Rui Wang, Robin Walters, and Rose Yu. Incorporating symmetry into deep dynamics models for improved generalization. arXiv preprint arXiv:2002.03061 , 2020.

Rui Wang, Robin Walters, and Rose Yu. Approximately equivariant networks for imperfectly symmetric dynamics. arXiv preprint arXiv:2201.11969 , 2022.

Maurice Weiler and Gabriele Cesa. General e (2)-equivariant steerable cnns. Advances in Neural Information Processing Systems , 32, 2019.

Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco Cohen. 3d steerable cnns: Learning rotationally equivariant features in volumetric data. arXiv preprint arXiv:1807.02547 , 2018.

Daniel E Worrall, Stephan J Garbin, Daniyar Turmukhambetov, and Gabriel J Brostow. Harmonic networks: Deep translation and rotation equivariance. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition , pages 5028-5037, 2017.

Jin Xu, Hyunjik Kim, Thomas Rainforth, and Yee Teh. Group equivariant subsampling. Advances in Neural Information Processing Systems , 34:5934-5946, 2021.

Jianke Yang, Robin Walters, Nima Dehmamy, and Rose Yu. Generative adversarial symmetry discovery. arXiv preprint arXiv:2302.00236 , 2023.

Raymond A Yeh, Yuan-Ting Hu, Mark Hasegawa-Johnson, and Alexander Schwing. Equivariance discovery by learned parameter-sharing. In International Conference on Artificial Intelligence and Statistics , pages 1527-1545. PMLR, 2022.

Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Ruslan Salakhutdinov, and Alexander Smola. Deep sets. arXiv preprint arXiv:1703.06114 , 2017.

Allan Zhou, Tom Knowles, and Chelsea Finn. Meta-learning symmetries by reparameterization. arXiv preprint arXiv:2007.02933 , 2020.

Hongpeng Zhou, Minghao Yang, Jun Wang, and Wei Pan. Bayesnas: A bayesian approach for neural architecture search. In International conference on machine learning , pages 7603-7613. PMLR, 2019.

Barret Zoph, Vijay Vasudevan, Jonathon Shlens, and Quoc V Le. Learning transferable architectures for scalable image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 8697-8710, 2018.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nEdit: Rating updated from 6 to 7 after rebuttal.\n\nThe goal of the paper is to learn an interpolation between non-equivariant and equivariant models. The authors introduce different convolutional and non-convolutional linear layers, optionally being sparsified via factorizations or a smooth spatial basis. The basic idea is to define the model as a linear combination of translation-equivariant and non-equivariant layers, and to optimize their relative contribution in order to select whether the final model is equivariant or not. This is achieved by posing different Gaussian priors on their parameters, whose widths constitute hyperparameters to be optimized via Bayesian model selection. This model selection differs from prior work, which tuned the hyperparameters explicitly.\n\nBayesian model selection requires the evaluation of a marginal likelihood term, which is infeasible. This is addressed via the Laplace approximation of the parameter posterior proposed by Immer et al. (2022). As this would still require the computation of a Hessian matrix which scales quadratically in the number of parameters, a Kronecker factorization approximation is used.\n\nExperiments on image classification datasets show that the interpolated models outperform non-equivariant models and, on tasks with broken equivariance, also strictly equivariant models. The latter are, in turn, performing slightly better on strictly equivariant learning tasks. An ablation to a pure MAP optimization baseline shows that the model selection approach results in significant gains.\nResults for further groups beyond translations are discussed in the appendix and briefly evaluated in the experimental section.\n\nSTRENGTHS:\nThe paper combines ideas from prior work and demonstrates a superior performance compared to these baselines. It relies on a Bayesian approach and shows how it can be made feasible via approximations despite the analytical intractability and high dimensionalities of parameter spaces. Automatically learning hyperparameters is certainly an improvement over manual tuning. The presented empirical evidence shows that the model selection does indeed select equivariant or non-equivariant layers depending on the symmetries of the learning task.\n\nWEAKNESSES:\nThe main downside of the approach is that it does not really learn equivariant models from scratch, but rather learns to select between pre-specified models with different levels of equivariance. However, this downside is shared by a line of prior work, which the current approach improves upon.\n\nI am also not entirely convinced that learning such selections is practically relevant since the appropriate levels of equivariance groups are usually known a-priori or become evident when comparing an equivariant model against a non-equivariant baseline. The experiments suggest that using strictly equivariant models is working better in cases where the desired equivariance group is indeed known a-priori.\n\nInstead of describing general equivariant mappings, the main paper considers only conventional convolutions, i.e. translation equivariance. More general results are discussed in the supplementary material and briefly evaluated in the last experimental section. It would have been nice to have a more general formulation in the main paper.\n\nIt would also have been interesting to see how the method scales to multiple symmetry groups at once. For instance, one could consider all 2^3 (i.e. exponentially many) combinations of using translations, rotations and reflections.\n\nThe explanation of the method could also be more clear. I had to re-read section 3 and the exact nature of the hyperparameters, introduced in section 3, remained vague till section 4.2. It would have been easier to just mention model selection of hyperparameters in the intro and moving the current section 3 between sections 4 and 5.\n\nQUESTIONS:\nThe authors state that \"adding symmetries does not directly improve training losses that rely on training fit\", however, many papers on equivariant models show their improved convergence rate and final loss when being fitted. Could this statement be clarified and supported more explicitly with evidence?\n\nThe notation of the limit in equation 9 does not seem to make sense. I guess that the authors just want to say that the implication follows given $\\sigma^2=0$?\nI found the notation of the feature map's domain as $\\mathbb{Z}^3$ somewhat confusing and would rather write $[0,C]\\times \\mathbb{Z}^2$ (spatially supported on $[0,X]\\times[0,Y]$).\nMore questions/suggestions are found in the \"weakness\" section.\n\nLIMITATIONS:\nLimitations are not explicitly discussed. I do not have ethical concerns", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-19"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nWhile (group) convolutions encode strict symmetries into neural network architectures, this paper presents a method for representing flexible symmetry constraints and learning the degree of symmetry automatically (through marginal likelihood objectives). Their method builds on residual pathways to represent each NN layer as the sum of fully connected and convolutional layers, where the initialization variance hyperparameter of the FC weights controls the degree of equivariance and is automatically optimized through the marginal likelihood. They introduce a number of techniques to make this process tractable in practice:\n* The FC layer is factored along spatial dimensions and possibly sparsified\n* To efficiently compute a laplace approximation of the marginal likelihood, they use a KFAC approximation of the Hessian. The factorization/sparsification of the FC layers admit additional simplifications to the KFAC Hessian computation.\n\nThe paper includes experiments comparing this method with plain FC or CNN architectures on synthetic and natural image datasets (MNIST and CIFAR-10), and show that it can learn to adjust the degree of equivariance to achieve good performance on these tasks. Analyzing the optimized hyperparameters shows that the method prefers to make earlier layers more equivariant and later layers less equivariant, in agreement with common architecture design principles. Finally, they show the ability to select between multiple symmetry groups (Conv and 90degree rotation GConv) on CIFAR10.\n\nSTRENGTHS:\nThe paper introduces a number of technical innovations to tackle a very challenging problem, automatic symmetry discovery. Marginal likelihood optimization removes the need for validation datasets or handcrafted regularizers when learning symmetries (which were used by some prior work), but is difficult to compute efficiently at the scale of modern NNs. It also attempts to address a common problem of symmetry discovery work where, if you start with no symmetry assumptions (FC), you end up with an enormous number of parameters for high dimensional inputs which is not computationally feasible in practice. The authors show that factoring (along spatial dimensions) and sparsifying the FC layers can reduce the number of parameters and even admit more efficient marginal likelihood estimation.\n\nWEAKNESSES:\nThe empirical results are limited and do not demonstrated that this method will be broadly applicable, in my opinion. The real datasets considered (CIFAR-10 and MNIST) are relatively simple and existing techniques (e.g., strong data augmentation or SSL + resnets or ViT) likely obtain much stronger results on these benchmarks without needing symmetry learning. Although the purported advantage of symmetry learning is that it can hopefully do better than human chosen symmetry constraints (like humans choosing data augmentations), the empirical results here don't show that.\n\nAlthough in principle the method is applicable to any symmetry group, the experiments seem to focus almost exclusively on translation invariance. Only 6.3 studies the case with 2 possible symmetries: translation (conv) and rotation (GConv). Is there a concern with scalability to more possible symmetries? Ideally, we would like the method to be able to learn any relaxed symmetries more generally without having to restrict the search space to one or two options (translation and rotation).\n\nRelatedly, much of the paper and in particular the spatial factorization seem particularly suited to image inputs (or similar input modalities). Would spatial factorization still work well with other data modalities, like graphs representing molecules? The experiments also seem to be focused on image inputs only.\n\nRelated work: the experiments and aims of this paper are reminiscent to [Elsayed et al, 2020](https://arxiv.org/abs/2002.02959), which also aimed to learn related spatial symmetry by operating on a spectrum between locally connected and convolutional layers. Using locally connected layers can also be viewed as \"sparsifying\" the fully connected layer to decrease computational costs. It would be interesting to discuss the differences and similarities here, either in text or in terms of empirical results (or both).\n\nQUESTIONS:\nSome minor things:\n\nWhat is P in Eq 1, the number of parameters? Might've missed this.\n\nWhy does the definition of fully connected preserve the spatial dimensions X and Y? In principle, I'd expect that a fully connected layer (or a conv layer) can change the spatial dimensions. Also, the closing bracket $[0,X$ is missing above Eq 2.\n\nSection 6.2: If I understand properly, isn't the optimal behavior just to ignore the fully connected pathway altogether and set its variance -> 0?\n\n[Post-rebuttal update]: I have read the author's response. They answer most of my questions. I largely maintain my original score and evaluation.\n\nLIMITATIONS:\nThe authors discuss limitations of their method and the general line of work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a neural network architecture and gradient-based training algorithm for modeling approximately equivariant functions. The architecture builds upon residual pathway (Finzi et al., 2021), where each layer of network is parameterized as an additive combination of constrained equivariant path and unconstrained fully connected path. A general challenge for such architectures is that equivariance is not favored for fitting training data, so empirical loss minimization will likely result in unstructured solutions that do not generalize. While the residual pathway paper solves this by putting higher weight regularization on unconstrained path, the strength of regularization is a hyperparameter that requires search with validation data in principle. The main motivation in this paper is to learn the extent of equivariance in each layer only from training data in a single training run (online). For this, the authors adopt Bayesian model selection which allows gradient-based learning of hyperparameters from training data by maximizing marginal likelihood estimate, specifically chosen in this work to be Laplace approximation with structured Hessian approximation with KFAC. Given that, the major technical contributions of the paper are on (1) improving the parameter efficiency of residual pathway by introducing convolution on Lie groups as well as factorization and spatial sparsification based on standard exponential basis functions, and (2) specifying the extent of equivariance as hyperparameters controlling the priors placed on the parameters, and (3) deriving KFAC for the proposed parameterizations so that gradient-based learning of hyperparameters controlling the extent of equivariance is made possible through maximization of marginal likelihood estimate. The authors provide experiments mainly regarding discrete 2-dimensional translation symmetry, and demonstrate that the proposed algorithm can (1) learn partial symmetry when needed, (2) recover the standard architecture of convolutional stack postfixed by fully-connected layers, and (3) learn to select from multiple symmetry groups depending on task, solely from training data.\n\nFinzi et al., Residual Pathway Priors for Soft Equivariance Constraints (2021)\n\nSTRENGTHS:\nS1. Overall, I think this is a solid work that contributes towards solving a challenging and important problem of learning symmetry from training data by bridging (approximate) equivariant architecture design and Bayesian deep learning. I find no critical issue with originality, quality, clarity, and significance of the work; the writing is overall clear, the design of the algorithm is clearly motivated and presented, and the experimental results seem to support the main claims and motivations.\n\nWEAKNESSES:\nW1. In Table 1-3, in addition to test performance, it would be nice if I could see how the models (over)fit to training data (i.e., how they generalize) given that a major motivation of the work comes from the fact that equivariance is not encouraged when fitting training data but beneficial for generalization.\n\nW2. The algorithm is empirically demonstrated for discrete 2-dimensional translation and 90-degree rotation symmetries. While the authors argue that extension to more general groups is possible in principle, I find the empirical demonstration regarding the argument is limited compared to the residual pathway paper (Finzi et al., 2021) that provided comprehensive experiments regarding e.g., continuous orthogonal groups as well.\n\nFinzi et al., Residual Pathway Priors for Soft Equivariance Constraints (2021)\n\nQUESTIONS:\nQ1. Considering the space of input and output feature maps, how is the group lifting (Appendix A) done in general when there is a set of M groups G1, …, GM (Appendix B)? Do you use direct product G = G1 x … x GM? This doesn’t seem straightforward, I might have missed something.\n\nQ2. Does the derivation of KFAC in Appendix C extend to other groups in a similar way to Appendix A? It seems so intuitively, but I would like to ask for a clarification from the authors.\n\nQ3. Do the layer indices 0-7 in Appendix D match each layer in Table 6? That is, can I regard the learned layer 6 in Appendix D to correspond to the layer that maps spatial dimension 4 x 4 to 2 x 2 in Table 6?\n\nQ4. In Table 6, it is written that the architecture is used for all experiments, and CONV layers are marked with kernel size (e.g., 3 x 3). How should one interpret this for S-CONV layers where, to my understanding, kernel size is not fixed?\n\nLIMITATIONS:\nThe authors have partially addressed limitations of the work in Section 7; I encourage further clarifications of limitations that needs to be addressed in future work, if there are any.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}

comments: SUMMARY:
The work proposed an automatic way to learn equivariance in each layer by finding a balance between the equivariant layer and the unrestricted fully connected layer. Unlike the previous work on soft equivariance, the work proposed to learn the balance between them via Bayesian model selection using gradients. The work also proposes different parameter-reduction techniques and achieves better results in the conducted experiments.

STRENGTHS:
1. The work provides a technique for learning equivariance structure automatically through gradient
2. The work addresses the overparameterization of such soft equivariant models by factorization and sparsifications.
3. The paper is well-organized and clearly written.

WEAKNESSES:
1. Evaluations: The empirical evaluations are conducted primarily on image data where the motivation for finding a balance between the equivariant and non-equivariant layers is not clear (except for the toy problem). Clearly, for image classifications, labels should not be affected by translation or rotation.  For the toy problem, a comparison with Finzi et al. [2021a] is not provided.
2. The performance of regular equivariant architecture is not provided, which makes it difficult to measure the gain by the proposed method.

Marc Finzi, Gregory Benton, and Andrew G Wilson. Residual pathway priors for soft equivariance
constraints. Advances in Neural Information Processing Systems

QUESTIONS:
1. I am curious about the memory and time (training) required for the proposed model compared to the regular equivariant architecture.
2. What may be the reason for MAP performing worse than the proposed Diff. Laplace method?

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-10