Learning Layer-wise Equivariances Automatically using Gradients
Authors: Tycho F.A. van der Ouderaa
Alexander Immer
Mark van der Wilk
Venue: NeurIPS 2023 spotlight (2023)
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
Tycho F.A. van der Ouderaa'
Alexander Immer
Mark van der WilkI4
Departmeet of Computing, Imperial College London, United Kingdom
DDepartment of Computer Science, ETH Zurich, Switzerland
MMa Planck Institute for Intelligent Systems, Tübingen, Germany
*Department of Computer Science, University of Oxford, United Kingdom
Abstract
Convolutions encode equivariance symmetries into neural networks leading to better generalisation performance. However, symmetries provide fixed hard con- straints 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 equívari- ances. Secondly, symmetries act constraints and 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 au- tomatically learn layer-wise equivariances on image classification tasks, achieving equivalent or improved performance over baselines with hard-coded symmetry.
Introduction
Symmetry constraints, such as layer-wise equivariances of convolutional layers, allow neural networks to generalise efficiently [Cohen and Welling, 2016]. However, it not always clear which symmetries to use and how much they should be enforced. In digit recognition, some rotational robustness 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 al., 2021a] and non-stationary filters [van der Ouderaa et al., 2022], effec- tively 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 encouraged by training losses that measure data fit, noted in 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 learn layer-wise equivariances from training data single training procedure.
We demonstrate automatically learning wise symmetry structure on image classification tasks. To do so, we improve upon existing parameterisations of differentiable equivariance and derive corresponding Cronecker-factorre classification, show that our method automatically learns convolutional structure in early layers and achieves similar or improved performance compared to architectures with hard-coded symmetry.
Laplace approximations to the marginal likelihood. On image
Code accompanying this paper available at httpa: h/tpa::///ttbbb wwt..cco//ttaeeee htt/ttyycooo/e//ee eella 37th Conference on Neural Information Processing Systems (NeurIPS 2023).

---

2 Related Work
Enforcing strict equivariance into the architecture. 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, 8,, 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, 8, Kosiorek et al., 2019], as formalised in Cohen and Welling [2014], Cohen et al. [2018]. Positional information under striet translation equivariance of convolutions can only be place-coded, not rate-coded in features. Capsule networks aim to break this strict disentanglement 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 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 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 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- stationnay 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, 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 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 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 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 convolutional architectural struc- ture [Cohen et al., 2019], symmetry discovery can be viewed as form of neural architecture search (NAS). Our approach uses relaxed symmetries to allow differentiable search over architectural struc- tures, similar to DARTS [L al., 2018a) 1a]. 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 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 - well-understood statistical procedure. Traditionally, model selection has only been available for models with tractable marginal likelihoods, such as Gaus- sian Processes [Rasmussen, 2003], but approximations for deep learning have also been developed, such as Variational Inference schemes [Ober and Aitchison, 2021], J. deep Gaussian Processes [Dutor- doir 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], 1]. but often use lower bounds that only work on shallow networks. For deep networks, Immer 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 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 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 map- pings 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 2- dimensional translation symmetry ZZ, 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| x + IR (spatially supported at intervals [0, X] x [0, YD, with C input channels outputting + IR with the same spatial support:
C' channel feature map y [0, x Z
y(d, (c, = Σ
Σ
x(c, ((' C,, 2, vf x.,) v)
C
J.
where non-zero weights in w: - + R can be represented in C'Cx²y2 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 Wa and write as a matrix W E X 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.
3

---

Convolutional layer in the input (dz, dy) E Z² yield equivalent translations in the output, mathematically:
We might want to constrain a layer to be strict equivariant, so that translations
y(d,a x' + dx, y + dg) = >Σ x(e, x + dr, dy) C, ,','''',, Vdr, dy E Z
(2)
2.8
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 F = x'
and y = y' y and can thus be written in terms of stationary weights (c, d, T. 9) + l::
y(d, t, ,/) = ΣΣ
x(c,,,,)0(c (d) e, - I, v/ y),
J.y
with C'CXY parameters in 9. Note that Eq. (3) is exactly a convolutional layer ('convolution is all you need', Theorem 3.1 of Cohen et [2019]. Constraining the stationary weights be only locally supported on rectangular domain ã. [- S-1], and O otherwise, gives rise to small (e.g. 3 x 3) convolutional filters and further reduces the required parameters to C'OS".
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 0 + 0 with fully-connected 0 (c' ,e,,c,t x', x,2) + R and stationary ((C,, C, U) + R with h=x' - x and ÿー y:
y(c,, .'.5')))))) -E x(c, 2(c, 0) c, x' F. v)
Σ Σ
x,c, x. 0 c, 2. -y)
z.y
fully-connected (FC)
convolutional layer (CONV)
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 0 + 0, 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 ""CX*+''CX" or C'CX*+C'CS" 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 depth- wise [Chollet, 2017] and group separability [Knigge et al., 2022], we propose to factor fully-connected layers between input and output dimensions (c', c, x' ,,','', xr, y) 0, (c', c, y 02(c',,, c, xr y). This greatly reduces the parameters, while fully maintaining the benefits of fully-connected structure from full spatial H XW input and output H' x W' to both input and outputs channels CxC'.
ΣΣ x(,, (cc,,yy00 y)d(c c, D xr v) M.
Σ Or(d,c, v/) Σ x(c, 2,, 02(c c,,,y)
(5)
factored fully-commeeted (F-FC)
The factorisation of Eq. (5) has several benefits. First, re- duces the number of parameters required for the expensive the fully- connected path from C'Cx²y² to CC' 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 composition of two linear operations summing over Ex,, and E., which we use to derive fronecker-factored approximations to the second-order curvature and enable the use of scalable I Laplace approximations. Lastly, the parameters 0 and 0₂ have a single spatial domain that lends themselves to further optional sparsification with basis functions.
BY C
C'xC
Il' Ip'
0
HxW
H W
Figure 1: Eq. (5) visualised as tensors.

---

W2
Spatial sparsification To further improve the parameter efficiency. we consider parameteris- ing spatial dimensions in a lower dimensional basis sğ ",,,,,,,, C,
211 (2z, 8" .)
0 n. y)
22
Figllre 2: Spatial sparsification s - illustrated.
Σ azo; ((x. " f7). 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 oz. ( = Uj exp(-w²(( y o] =))²) for convolutional filters, with anchor point locations 2j E IR² and anchor point values uj E R. The work demonstrated that lengthscales w can allow differentiable control over equivariances (used in Section 3.2) well as practical sparsification with negligible performance loss, even if filters are already small (e.g. x 3).
ΣΣ> J.Y
x(c F. v)(() e, -,, v)
-EX
(cc,,,)) (c c, n v)
sparse convolution (S-CONV)
This reduces the number of parameters from C"CS² to C'CP with P anchor points, or C'CP + 2P if anchor point locations {z, 1MM (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 sf and s for spatially sparsified 02 and O2: x(c, a, v) (d, c, x, y)
Σ 0 (c' ,,,,,.v)
Σ x(c, x, () = Σ 8 (c', e, ,',y') Σ 20.3
2,9
spurse fully connected (S FC)
reducing the parameter cost from 2C'CXY to C'CP with P total number of anchor points. Com- pared to convolutional filters, which already have small spatial supports, sparsification of fully- connected layers will typically result in 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 e on groups, as explored in [Azangulov et al., 2022].
3.2 Specitying 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 (0, of I) and treat variances as hyperparameters n- [01, F2 .J.. network, such that the prior in each layer I enforces an equivariance symmetry in the following limit:
... We carefully parameterise the
= =
strict equivariance (Eq. (2)) and of > = 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 and (0 of I) with scalar variances of of e R and identity matrix T @ x (e) Intuitively, all weights in non-equivariant paths become zero in the zero variance limit =0, resulting in strict equivariance. We consider the same prior for the factored parameterisation F- FCC+CONV Likewise, for sparsified layers, we can place the prior directly on anchor points (N 0, of I) and (, of which will control equivariance through a similar mechanism. For S-CONV, this prior corresponds to the correlated prior of Fortuin et al. [2021] anchor points are placed exactly on the filter grid, whereas our anchor points are sparse M 5 and have learnable positions 2j. Alternatively, for sparsified S- layers, we can consider placing the prior directly on the lengthscales Nww 0, of and (wa 00, of 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 = forces wi to be zero, causing non-stationary components to be constant and weights to act stationary, resulting in strict equivariance. Treating lengthscales w as learnable hyperparameter has the additional advantage that it allows learning of convolutional filter frequencies from data through model selection of hyperparameter o7.
both residual pathways FC+CONV (0 of
Different constructions of the prior may induce different prior probabilities in function space. Nev- ertheless, 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 a==0 limiting case and relaxed layer-wise equivariance at higher variances 0j 00>0. Unlike prior works [Finzi et al., 2021a, van der Ouderaa et al., 2022] that require setting or tuning the prior variances of controlling equivariance constraints, we follow Bayesian model selection to learn them automatically from training data.
5

---

Bayesian model selection of symmetry constraints
Inspired by Bayesian methodology. we propose to empirically learn the amount of layer-wise equivariance in neural network from training data through approximate Bayesian model selection. define equivariances in our prior and treat the amount of symmetry as hyperparameterss optimised using the marginal likelihood well-understood statistical procedure known as 'empirical Bayes'. Mathematically, we find a point estimate of the hyperparameters n = arg max, p(D(m) whilst integrating out the P model parameters 0 E IR":
((D = p(100 ())
The marginal likelihood p(D)n) forms the normalising constant of the unnormalised posterior over parameters:
p n) ((D)O,m)p(0) p(D(n)
(DD m)p(0)
Jo P(D ())(
As the marginal likelihood is constant in 0, it can be ignored when training neural network parameters with regular MAP point-estimate arg ming Lo:
Lg log PD) 0, n) log p(0)
To optimise hyperparameters, the marginal likelihood p(D(n) 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 μ O. and the covariance set to the local curvature Σ = H = -- V2 that point:
P n) g(B)D. NH,
Integrals of Gaussians are known in closed-form. and thus the normalising constant of to estimate the log marginal likelihood objective used to optimise hyperparameters 77,- = arg min, Ly
Cn = -loo (v )2x)² Ho. p(D, 0.1m)) = log PDD)O. - log.llog))))) p(O.) - log(( (T)) + log He.
NLL Data fit
Occam's
(9)
Although the Laplace approximation greatly simplifies the marginal likelihood, the Hessian deter- minant He. remains difficult to compute for large neural networks due to quadratic scaling in parameter count (( In practice, we therefore use local linearisation [Bottou et al., 2018] re- sulting 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 form of Occam's '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 Co w.r.t. data augmentation parameters prefers collapse into solution without symmetry constraints resulting in fully- 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-awan minimisation and has recently been used to learn invariances large ResNets [Immer et al., 2022].
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

---

(110))), the loss becomes intractable for large networks. We use KFAC [Martens and Grosse, 2015] approximate the Hessian of the log likelihood Ho. and have the Gauss-Newton approximation H:
Ho. H - H- En Hn, - En. J (a.) (xx))
where the Jacobians of network outputs f with respect parameters for data points æ,, are [J (x kk.pp Dfc. uated at 0.. KFAC is block- -oack-aa approximation of the GGN, with individual blocks = En Hn,:=En J(()))) ( for each layer I with parameters 0₁ and Jacobian J((x).
80, (x,,), and [A [a(x,)l..-- =
Plo ply, (.).fi..:0) is the log likelihood Hessian, all eval- 0fr0f,
For fully-come colly-connectd layer, Or with G and wER?!, data point w.r.t. parameters as Jı(æn) -a & It,n, where at,n ERR Dix1 is the input to the layer and gl.n ER. GR.. x is the transposed Jacobian of f w.r.t the 7th layer output. We thus have:
write Jacobians of single
H1- En J((xn) ) (I = Enlarn lasn @ ((nn) at, X gt.nl"
[,n @ Ig.. A A ( ~ N En at aTm @ [En A
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 learn layer-wise equivariances from data. The essential step to extend KFAC is to write the Jacobian of the Ith layer, (x), as Kronecker product allowing data points 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.
and S-CONV layers of Eqs. (6) and (7). allowing use in scalable Laplace approximations
5.1 KFAC for factored layers
To obtain the KFAC of F-FC in Eq. (5), we write the application of both weight parameters 0₁ and O₂ as matrix mutiplication akin to linear layer, which allows us to apply similar approximation as in standard KFAC. In particular, we first split up the operations in Eq. (5) into first applying 02 and then O,. Each of these operations can be written as matrix product @x with corresponding appropriate sizes of 0 and æ. In comparison to the standard KFAC case, we have Jacobian that can be written as
JØ = E. ? [y. = a. s Ic:
which means an additional sum needs to be handled. In particular, plugging the Jacobian structure into the standard KFAC approximation, we have
Hg 23 Σn, Jo(xn) A(Ix)) = =n an. & A(x,) an, @ gn,e]
= =En,c,e ( n. 8 (9 (xn))gnee NC I'm an n,el 3 IEn,c In.e A(x Do.t.
where the first approximation is from the Hessian to the GGN and the second approximation is due KFACss 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 or values. In App. C, we derive the KFAC for 0 and 0₂ in detail.
5.2 KFAC for sparsified layers
To obtain KFAC for sparsified S-FC and S-CONV layers, 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 our new anchor points parameters, instead of the induced weights st we apply the chain rule:
0f Öf as Da+ 0s da,"
d (d, Dag
using the partial derivative
"x v)
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.
7

---

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 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 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 the following sections. We slightly altered the MNIST dataset [LeCun et al., 1989] so that digits are randomly placed 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 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 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.
Example:
MAP
Diff. Laplace
FC
CONV
FC
CONV
Learned (ours)
3
Symmetry
Prediction task
112.7 M
0.4 M
112.7 M
0.4 M
1.8 M
Strict symmetry
(digit)
95.73
99.38
97.04
99.23
98.86
(3, upper left)
Partial symmetry
(digit, quadrant)
95.10
24.68
95.46
24.50
99.00
Table 1: Preventing symmetry misspecification. Test accuracy for non-symmetric FC, strictly symmetric CONV, and learnable symmetry F-FC+CONV models on translation invariant task and task that can not be solved under strict translation symmetry. Unlike the FC and CONV baselines, the proposed model with learnable symmetry constraints achieves high test performance on both tasks.
6.2 Learning use layer-wise equivariant convolutions on CIFAR-10
Convolutional layers provide useful inductive bias for image classification tasks, such 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 the architecture of Neyshabur [2020] (see App. F). In Table 2, we compare test performance when using fully-connnctted 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 0, anchor points u, or lengthscales w, 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 x)) performance as the convolutional CONV models.
Learaalee
Prior
MAP
Difl. Laplace (Ours) Test Acc. (?)
Layer
equivariance
# Params
Test NI. (L)
Test Acc.
Test NL).
C,(())
FC
0
133.6 M
2.423
64.42
1.319
$2.36
1.896
CONV
82.81
84.07
1.022
FC+CONV
133.9 M
1.713
76.93
0.489
83.32
1.019
F-FC
4.291
50.47
1.307
53.38
2.119
C.FC (Ours)
0
188 M
1.343
81.61
0.468
83.92
1.277
S.FC
14
0.8
5.614
50.21
1.496
46.18
1.926
S-CONV
0.2 M
0.604
82,53
0.562
81.29
1.037
S-FC+S-CONV (Ours)
M
5.485
$8.80
0.538
81.55
0.978
S-FC
7.783
49,69
1.319
$2.22
1.813
S-CONV
0.2
0.913
81.84
0.549
81.20
1.055
S-FC+S-CONV (Ours)
M
2,040
77.97
80.63
1.067
Table 2: Equivariant convolutional layers CONV obtain the highest performance on image classi-
fication task when trained with regular MAP. When maximising Laplace approximated marginal likelihood, models with learnable symmetry learn from training data to 'become convolutional" (see Section 6.2 for a discussion) and obtain equivalent (within <1%) performance as CONV models.
8

---

Layers learn to become convolutional, except at the end of the network. Upon further inspection of learned symmetries, we observe very high prior precisions of >>>1>00 10 for non-equivariant paths in most layers, thus negligibly contributing to final outputs (see App. D.I). 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 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 coincides with common practice in literature. Spatial dimensions in the used architecture are halved through pooling operations, reaching with sufficient or global pooling, after which features become strictly invariant and there is no difference anymore between Ixl 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 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 per-layer basis.
Equiivarance frest inest layers
32
CONV
CONV
No equiv faancce the end of modill
CONV
CONV
8x8
CONV
CONV
4x4
Finall Finallayers respect symmetry, because
of spatial dims
CONV
FINAL
FINAL
Figure 3: Learned CONV+FC model.
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+-CONV 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 (a=10 to a=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. and effective number of parameters (App. D.2), we observe a positive correlation between learned symmetries and augmentations applied to the datasets trained on.
with Diffrectishic Laphaz
Elfcative Num.
Paraus NLL suxunay
Test ())
sasunacy (T)
Appecs, Maggiik (
FC
CONV OCONV (8)
MNIST
1.2M
0.172
97.5g
0.033
99.21
0,328
10
(0-46)
15
(0-98)
15
(1-100)
Teamlated MNIST
1.2M
0.812
90.78
0.053
18.27
0,216
(0-0)
23
(0-9%
77
(1-10m
Rotated MNISS
1.2M
0.819
91.02
0.136
95.55
0.896
(0-20)
8
(0-47
83
(47-100)
CIFAR-10
1.2M
3.540
68.33
0.553
WI.54
0,926
(0-1)
(0-99)
56
(0-10m)
Rouxed CIAR
1.2M
5.953
45,30
1.266
55.68
1,30
(0:22)
(0 (0:41)
K2
(58 99)
Table 3: Selecting from multiple symmetry groups Negative log likelihood (NNL) and Laplace
learned symmetries measured by mean (min-max) relative effective number of parameters over layers.
7 Discussion and Conclusion
This work proposes a method to automatically leamn layer-wise equivariances in deep learning using gradients from training data. This 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 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 Cronecker-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 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 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. 249222200..9000 2022.
Erik Bekkers. B-spline cnns on lie groups. arXiv preprint arXiv: /2057, 2019.
Gregory Benton, Mare Finzi, Pavel Izmailov, and Andrew Gordon Wilson. Learning invariances in neural networks. arXiv preprint arXiv:2010. 1882, 2020.
Léon Bottou, Frank Curtis, and Jorge Nocedal. Optimization methods for large-scale machine learning. SIAM review, 60(2):223-31,, 2018.
Anadi Chaman and Ivan Dokmanic. Truly shift-invariant convolutional neural networks. In Proceed- ings 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. Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1251-1258, 2017.
Taco Cohen and Max Welling. Learning the irreducible representations 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 learming, pages 2990-2999, PMLR, 2016.
Taco S Cohen, Mario Geiger, and Maurice Weiler. Intertwiners between induced representations (with applications to the theory equivariant neural networks). arXiv preprint arXiv:1803.1074 2018.
Taco S Cohen, Mario Geiger, and Maurice Weiler. A general theory of equivariant cnns on homoge- neous 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: 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-22 2511, 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
PMLR, 2020.
Mare Finzi, Gregory Benton, and Andrew Wilson. Residual pathway priors for soft equivariance constraints. Advances in Neural Information Processing Systems, 34, 2021a.
Mare Finzi, Max Welling, and Andrew Gordon Wilson. A practical method for constructing equivari- ant multilayer perceptrons for arbitrary matrix groups. In International Conference on Machine Learning, pages 3318 3328. PMLR, 2021b.
Vincent Fortuin, Adrià Garriga-Alona 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 fronecker-factoredd approximate fisher matrix for convolution layers. International Conference on Machine Learning, pages 573-582. PMLR, 2016.
Geoffrey E Hinton, Sara Sabour, and Nicholas Frosst. Matrix capsules with em routing. In Interna- tional conference learning representations, 2018.
10

---

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 A. van der Ouderaa, Vincent Fortuin, Gunnar Rätsch, and Mark 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 Tch, 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/nmist// 1998.
Yann LeCun, Bernhard Boser, John Denker, Donnie Henderson, Richard Howard, Wayne Hubbard, and Lawrence Jackel. Handwritten digit recognition with a back pack-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 IEFE, 86(11):2278-25 2224,, 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: 00247,, 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 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, 039983, 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-1000 1994.
Kaitlin Maile, Dennis George Wilson, and Patrick Forré Equivariance-aware architectural optimiza- tion 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, ((():5776-58511 2020.
James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curvature. In International conference on machine learning, puges 1408-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 05:80788888,, 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 XrXiv:21 10. 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 62.62.226 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.07177, 2022.
Elise Van der Pol, Daniel Worrall, Herke van Hoof. Frans Olichoek, 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.05500 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-MICAAI 2018: 2 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 arXi::2002.03061, 2020.
Rui Wang, Robin Walters, and Rose Yu. Approximately equivariant networks for imperfectly symmetric dynamics. arXiv preprint arXi::201..11969, 2022.
Maurice Weiler and Gabriele Cesa. General e (2)-equivariant steerable cnns. Advances in Neural Information Processing Systems, 32, 2019.
12

---

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.02557. 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-500 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 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 1rXiv:7033.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 Quoe Le. Learning transferable architectures for scalable image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 8697-8710, 10, 2018.
13


Reviews:
[{"review_id": "pCeAcw0hNs", "reviewer": "Reviewer_iYmm", "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. ", "questions": "1. I am curious about the memory and time (training) required for the proposed model compared to the regular equivariant architecture.\n2. What may be the reason for MAP performing worse than the proposed Diff. Laplace method?", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The work provides a technique for learning equivariance structure automatically through gradient\n2. The work addresses the overparameterization of such soft equivariant models by factorization and sparsifications.\n3. 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.\n2. The performance of regular equivariant architecture is not provided, which makes it difficult to measure the gain by the proposed method.\n\nMarc Finzi, Gregory Benton, and Andrew G Wilson. Residual pathway priors for soft equivariance\nconstraints. Advances in Neural Information Processing Systems", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feedback.\n\n> motivation for finding equivariance balance \n\nSymmetry discovery from data is a well-recognized task and is actively studied [1,2,5,6,7]. The key motivation is that the true symmetry may be unknown, not specified, or that the optimal balance of equivariant and non-equivariant layers is in fact not 'clear'. Examples of this already arise in simple tasks, such as 6’s and 9’s in MNIST or our translational dependence toy problem. Finding the right balance arguably becomes even more difficult on more complex data, architectures and when optimal symmetries differ per layer, which we focus on. As discussed in Sec. 6.2, common vision architectures deploy translational symmetries, but are often not strictly invariant globally due to fully connected layers at the end of the network. Our method paves the way to discovering such favourable architectures automatically from training data.\n\nOur work demonstrates the principle of automatically learning layer-wise equivariances through approximate Bayesian model selection for neural networks. Even where strict symmetry is most desirable, our method is able to ignore non-equivariant layers and `become equivariant’, whereas standard MAP objectives cannot.\n\n> performance of regular equivariant architecture is not provided\n\nAll experiments include the performance of regular equivariant architectures denoted by 'CONV'.\n\n> memory and time (training)\n\nThe proposed factorised F- and sparsified S- parameterisations improve both training time and memory easily by 2x at a negligible loss in performance. Tables 1-2 include the exact amounts of memory in terms of the number of parameters and notice active GPU in practice to be roughly proportional to these counts. All experiments can run on a single 24GB GPU. We will report the exact numbers in the final manuscript.\n\nThe Diff. Laplace objective performs better than MAP but is ~1.2-2x slower in training. We would like to mention that we rely on modern linearised Laplace approximations, which are currently actively being studied and improved. For instance, concurrent work [4] shows a 10x improvement of the Laplace approximation and should also readily be applicable to our setting.\n\n> MAP performing worse than the proposed Diff. Laplace method?\n\nAlthough MAP can very successfully fit model parameters (θ), it can typically not learn hyper-parameters (η), such as prior variances or symmetry constraints (in our case: layer-wise equivariances). The reason that such a data fitting objective does not encourage symmetry is that symmetries constrain the functions that a network can represent. Consequently, the regular objective that maximises train data fit will always prefer as little as possible equivariance (no symmetry, so σ=0), even when more equivariance would be preferable in terms of better generalisation on validation/test data. This is why in Deep Learning it is common to select hyper-parameters using validation data, e.g. cross-validation, or more complex with RL outer loops [3] or differentiated validation losses [8].\n\nRecent work by [5] shows that the marginal likelihood can be used to learn invariances directly on training data and demonstrated in [6] for DL. Unlike regular MAP, this objective has a built-in 'Occam's razor' effect and therefore balances good 'data fit' with model complexity (through the log determinant term). Similarly, we use Diff. Laplace to perform approximate Bayesian model selection. Crucially, we introduce a scalable parameterisation that allows differentiable layer-wise equivariances, unlike prior works that only consider invariances. We demonstrate that - in line with theory - Diff. Laplace allows us to optimise the right layer-wise equivariances η using training data. This demonstrates the principle of automatically learning layer-wise equivariances from training data. \n\nWe hope that this shows that problems with the MAP objective are well-established facts. A more thorough discussion can be found in [5] and [6] and explicit examples that MAP cannot learn symmetry constraints in Sec. 4.2 of [7] and App. C of [6].\n\n[1] Benton, Gregory, et al. \"Learning invariances in neural networks from training data.\" NeurIPS 2020\n\n[2] Wang, Rui, et al. \"Approximately equivariant networks for imperfectly symmetric dynamics.\" ICML 2022\n\n[3] Cubuk, Ekin D., et al. \"Autoaugment: Learning augmentation strategies from data.\" CVPR 2019\n\n[4] Immer, Alexander, et al. \"Stochastic marginal likelihood gradients using neural tangent kernels.\" ICML 2023\n\n[5] Wilk, Mark van der, et al. \"Learning invariances using the marginal likelihood.\" NeurIPS 2018\n\n[6] Immer, Alexander, et al. \"Invariance learning in deep neural networks with differentiable Laplace approximations.\" NeurIPS 2022\n\n[7] van der Ouderaa, Tycho FA, et al. \"Learning invariant weights in neural networks.\" UAI 2022\n\n[8] Lorraine, Jonathan, et al. \"Optimizing millions of hyperparameters by implicit differentiation.\" AISTATS 2020\n"}, {"review_id": "LCNd4DnB1C", "reviewer": "Reviewer_gE9j", "summary": "This 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)", "questions": "Q1. 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?", "limitations": "The 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.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "S1. 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.", "weaknesses": "W1. 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)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback and help to improve the paper.\n\n> 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.\n\nThank you for raising this. We agree with the author that providing training performance is important given the motivation of learning symmetries by balancing data fit and model complexity through approximate Bayesian model selection. We can confirm that all models fit the training data well, as measured by a low negative log likelihood. We will add performance on the train set to the final version.\n\n> the authors argue that extension to more general groups is possible in principle, I find the empirical demonstration regarding the argument is limited\n\nThe method is described for general groups and an experiment demonstrating symmetry selection between multiple groups is given in Sec. 6.3. Although some prior works have considered a more extensive range of groups (e.g. [2]), our work has a much more extensive focus on the scalability of parameterisations and we put a considerable amount of work in deriving an objective function that can learn symmetry, which is a focus rather than how to parameterise to many groups. We demonstrate the principle of automatically selecting the most relevant layer-wise equivariances from training data through Bayesian model selection. We agree that it would be interesting to consider more groups in future work.\n\n> Q1. How is group lifting done in general?\n\nParameterising neural networks to obey group symmetries is an entire research field in its own right. The choice of parameterisation of the symmetry and its representation may vary depending on task-specific constraints. The formalisation of group lifting can be found in [3]. When combining rotation and translation, we follow regular representations similar to [1]. This means that if the network learns to ignore non-equivariant paths entirely, it will become exactly equivalent to standard G-CNN (e.g. P4CNN in [1]).\n\n> Q2. Does the derivation of KFAC in Appendix C extend to other groups in a similar way to Appendix A?\n\nIn our case, yes, since we define filters on a group through base features defined in a vector space [4] and can therefore readily be extended to KFAC without issues. KFAC only needs to properly handle the weight-sharing induced by the group parameterisation, which we explain how to do in Sec. 5. We do note, however, that parameterising symmetries in neural networks is a research field in its own right. Different parameterisations might affect how easily KFAC can be derived. We will add this disclaimer.\n\n> Q3. Do the layer indices 0-7 in Appendix D match each layer in Table 6?\n\nYes. The same architecture is used for all experiments and layer numberings are consistent.\n\n> Q4. 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\nThe kernel size of an S-CONV layer is fixed and kept the same (e.g. a 3 x 3 filter) as the CONV layer. The difference between both layers is that the number of parameters is decoupled from the kernel size. That is, by using basis features we can build an S-CONV layer with a 3 x 3 filter with fewer (e.g. 4) parameters, instead of the 9 required by CONV. A more extensive explanation of sparsifying filters through basis features can be found in [4]. We will make this more clear from the main text.\n\n\n[1] Cohen, Taco, and Max Welling. \"Group equivariant convolutional networks.\" International conference on machine learning. PMLR, 2016.\n\n[2] Finzi, Marc, Gregory Benton, and Andrew G. Wilson. \"Residual pathway priors for soft equivariance constraints.\" Advances in Neural Information Processing Systems 34 (2021): 30037-30049.\n\n[3] Kondor, Risi, and Shubhendu Trivedi. \"On the generalization of equivariance and convolution in neural networks to the action of compact groups.\" International Conference on Machine Learning. PMLR, 2018.\n\n[4] van der Ouderaa, Tycho FA, and Mark van der Wilk. \"Sparse Convolutions on Lie Groups.\" NeurIPS Workshop on Symmetry and Geometry in Neural Representations. PMLR, 2023.\n"}, {"review_id": "YEQ8syiZdU", "reviewer": "Reviewer_VChg", "summary": "While (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.", "questions": "Some 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.", "limitations": "The authors discuss limitations of their method and the general line of work.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The 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.", "weaknesses": "The 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback and help to improve the paper.\n\n> empirical results are limited\n\nEven when strict symmetries are desirable, our method shows that such favourable architectures can be discovered automatically, whereas the standard MAP objective cannot. Our proposed method allows automatically learning layer-wise equivariances from training data through approximate Bayesian model selection. We have demonstrated this principle but agree with the reviewers that assessing performance on larger real-world datasets and models would be interesting. \n\n> method is applicable to any symmetry group, but experiments focus on translation invariance\n\nWe describe the method for general groups and demonstrate symmetry selection between multiple groups in Sec. 6.3. Although some prior works have considered a more extensive range of groups (e.g. [1]), we have a greater focus on practical scalability of the parameterisations (Sec. 4) and puts a considerable amount of work in the objective function (Sec. 3, 5). In doing so, we have demonstrated the principle of automatically selecting the most relevant layer-wise equivariances from training data through Bayesian model selection. We agree that it would be interesting to consider more groups in future work.\n\n> scalability to more possible symmetries?\n\nScalability to more symmetries likely depends on the exact problem set-up and considered groups and representations. Parameterising group symmetries is a research field in its own right and not the focus of this work. Nevertheless, the existing parameterisation of neural networks with many symmetries and regular group representations should be readily extendable to relaxed and learnable symmetries using the techniques presented in this paper. We see no direct concern with scalability to more symmetries. This is an interesting avenue for future research, but orthogonal to enabling learnable layer-wise symmetries from training data, which this work focuses on.\n\n> Would spatial factorisation still work well with other data modalities, like graphs representing molecules?\n\nYes, the spatial factorisation should generalise to other data modalities as long as you can construct basis functions on the manifolds [2], which are also available for graphs [4].\n\n> Related work: the experiments and aims of this paper are reminiscent to Elsayed et al, 2020, \n\nWe greatly thank the reviewer for pointing this out and will include [3] in the related work.\n\nThere are indeed some interesting parallels between the work and the parameterisation aspects (Sec. 4) of our work. Indeed, the motivation behind including spatial dependence is similar as well as recognising that \"combining weights scales linearly with the number of pixels in the image, which may be large.\" There seem to be subtle differences in how both methods address this. For instance, [3] uses 'low-rank locally connected layers' factorisation between height and width dimension, whereas our factorisation factors between the full input and output dimensions. Further, our sparsification uses basis features which do not necessarily result in a rank-1 matrix spatially, whereas the low-rank factorisation of [3] will.\n\n> What is P in Eq 1, the number of parameters?\n\nIndeed, P is the number of parameters. We will make this clear from the text.\n\n> Why does the definition of fully connected preserve the spatial dimensions X and Y?\n\nThis indeed differs slightly from the typical formulation of fully-connected layers that allow different input and output dimensions. We do this mainly for notational simplicity and to avoid running into the complexity of having different group structures on the input and output. We are still absolutely flexible regarding changing spatial dimensions but only do this through equivariant pooling (see App. F). For instance, instead of defining an FC layer with 2x smaller output dimensions, we simply use an FC without spatial downsampling followed by a 2x2 filter spatial pooling layer. We go through all this effort to ensure that equivariant paths are strictly equivariant and that our experimental set-up can verifiably show that our method actually learns layer-wise symmetries. This might be less important when applying the method in practice.\n\n> If I understand properly, isn't the optimal behaviour just to ignore the fully connected pathway altogether and set its variance -> 0?\n\nYes. Equivariance on the image classification task is likely desirable, thus ignoring the fully connected pathway. We refrain from making any hard claims about optimality but generally agree with the reviewer.\n\nIn line with theory, we observe that this “optimal behavior” cannot be learned with the standard MAP objective but can be learned with Diff. Laplace. Using the proposed objective, the network `becomes equivariant’ and largely learns to ignore the fully connected pathways. This demonstrates the ability of our method to select the most relevant symmetry from training data automatically.\n\n[1] Finzi, Marc, Gregory Benton, and Andrew G. Wilson. \"Residual pathway priors for soft equivariance constraints.\" NeurIPS 2021\n\n[2] Azangulov, Iskander, et al. \"Stationary kernels and Gaussian processes on Lie groups and their homogeneous spaces i: the compact case.\" 2022\n\n[3] Elsayed, Gamaleldin, et al. \"Revisiting spatial invariance with low-rank local connectivity.\" ICML 2020.\n\n[4] Borovitskiy, Viacheslav, et al. \"Matérn Gaussian processes on graphs.\" AISTATS 2021.\n\nTypos will be fixed.\n"}, {"review_id": "jRQUZpzvcc", "reviewer": "Reviewer_n5aj", "summary": "Edit: 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.", "questions": "The 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.", "limitations": "Limitations are not explicitly discussed. I do not have ethical concerns", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The 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.", "weaknesses": "The 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback and help to improve the paper.\n\n> learns to select between pre-specified models with different levels of equivariance. However, this downside is shared by a line of prior work + scales to multiple symmetry groups at once\n\nThe approach is general in that it can work with general groups, as long as they can be differentiably parameterised. Learning symmetries using training data only from simple affine transformations is still actively investigated [1, 2, 3], and we extend the parameterisations to allow for layer-wise equivariances and provide concrete improvements to the scalability of parameterisations. Parameterising distributions over other groups and scaling to multiple symmetry groups is an interesting avenue for future work, but orthogonal to enabling differentiably learning of such parameterisations, which is what our paper focuses on.\n\n> using strictly equivariant models is working better in cases where the desired equivariance group is indeed known\n\nWe agree with the reviewer that when a strict global symmetry is known a-priori or self-evident, it would always be better just to build it into the model. However, in some cases, the symmetry might not be known, specified, or should not be strictly enforced (think 6’s and 9’s in MNIST under rotational symmetry). Finding the right symmetry or optimal balance between equivariant and non-equivariant layers becomes harder on more complex models and datasets, especially when this can also differ layer-wise.\n\nOn image classification tasks where equivariance seems very desirable, the standard MAP objective cannot select and ignore the non-equivariant pathways. Yet, our proposed Diff Laplace objective based on approximate Bayesian model selection does allow symmetry learning and can automatically select the most relevant symmetry, becoming largely equivariant on this task. This demonstrates the ability to learn relevant layer-wise symmetries from training data automatically.\n\n>discussed in the supplementary material and briefly evaluated in the last experimental section\n\nWe will follow the suggestion by the reviewer to move the description of general group equivariant mappings to the main text.\n\n> moving the current section 3 between sections 4 and 5.\n\nWe agree with the reviewer and will adopt the suggestion to move Sec. 3 between Sec. 4 and Sec. 5 to make the paper clearer.\n\n> The 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\nThere is significant evidence that MAP alone cannot learn hyperparameters, including symmetry constraints. This is supported by experiments performed in earlier work on invariance/equivariance learning [1,2,3], with explicit examples that MAP cannot learn symmetry constraints in Sec. 4.2 of [2] and App. C of [3]. Because of this reason, works in literature often require differentiated validation losses [5, 6], explicit regularisation [7], or RL outer loops [4]. We hope this shows that problems with the MAP objective are well-established facts. \n\nThe reason that such a data-fitting objective does not encourage symmetry is that symmetries constrain the functions that a network can represent. Consequently, the regular objective that maximises train data fit will always prefer as little as possible equivariance (no symmetry, so σ=0), even when more equivariance would be preferable in terms of better generalisation on validation/test data.\nIn some particular cases, as the reviewer mentions, it can happen that there is an improved training fit when using the right equivariance. However, this only occurs when model capacity is constrained. To properly compare two architectures, they need to be large enough for their performance not to be artificially restrained by a lack of size. This is where the difference in training fit disappears and where an additional measure of complexity needs to be used, which the Laplace approximation provides.\n\nWe have seen similar claims about faster convergence rates in models with the right equivariance, and we share the reviewer's view that this is a phenomenon that _may_ allow point optimisation to favour architectures that generalise well. However, there are currently no papers that successfully demonstrate using convergence rates for architecture selection, and our attempts to use this have failed, similarly to [1,2,3].\n\nWe show that by using the proposed Diff Laplace objective using approximate Bayesian model selection, we can in fact automatically learn symmetries using training data.\n\n> notation of the limit in equation 9\n\nIndeed. Eq.9 should be the implication that follows given $\\sigma^2 = 0$. We will fix this.\n\n> I found the notation of the feature map's domain as somewhat confusing \n\nWe thank the reviewer for pointing this out and adopting the suggested change in notation.\n\n\n[1] Wilk, Mark van der, et al. \"Learning invariances using the marginal likelihood.\" NeurIPS 2018\n\n[2] van der Ouderaa, Tycho FA, et al. \"Learning invariant weights in neural networks.\" UAI 2022\n\n[3] Immer, Alexander, et al. \"Invariance learning in deep neural networks with differentiable Laplace approximations.\" NeurIPS 2022 \n\n[4] Cubuk, Ekin D., et al. \"Autoaugment: Learning augmentation strategies from data.\" CVPR 2019\n\n[5] Lorraine, Jonathan, et al. \"Optimizing millions of hyperparameters by implicit differentiation.\" AISTATS 2020\n\n[6] Liu, Hanxiao, et al. \"Darts: Differentiable architecture search.\" arXiv preprint arXiv:1806.09055 (2018).\n\n[7] Benton, Gregory, et al. \"Learning invariances in neural networks from training data.\" NeurIPS 2020\n\nTypos will be fixed."}]

Author rebuttal:

Meta-review: All reviewers support the acceptance of this paper and the AC agrees with this evaluation. Overall, it is a well motivated paper on a challenging problem. A side note on additional references [A,B] that are very related to this work.

[A] Yeh, Raymond A., et al. "Equivariance discovery by learned parameter-sharing." International Conference on Artificial Intelligence and Statistics. PMLR, 2022.

[B] Yang, Jianke et al. “Generative Adversarial Symmetry Discovery.” International Conference on Machine Learning (ICML), 2023
Decision: Accept (spotlight)