SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $\mathrm{FC + CONV}\mathcal{N}(\pmb {\theta}|\mathbf{0},\sigma_l^2\pmb {I})$  and  $\mathcal{N}(\bar{\theta} |\mathbf{0},\bar{\sigma}_l^2\mathbf{I})$  with scalar variances  $\sigma_l^2,\bar{\sigma}_l^2\in \mathbb{R}$  and identity matrix  $I\in \mathbb{R}^{|\pmb {\theta}|\times |\pmb {\theta}|}$ . Intuitively, all weights in non-equivariant paths become zero in the zero variance limit  $\sigma_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  $\mathcal{N}(\pmb {u}|\mathbf{0},\sigma_l^2\pmb {I})$  and  $\mathcal{N}(\bar{\pmb{u}} |\mathbf{0},\bar{\sigma}_l^2\mathbf{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  $\mathcal{N}(\omega |\mathbf{0},\sigma_l^2\mathbf{I})$  and  $\mathcal{N}(\bar{\omega}_l|\mathbf{0},\bar{\sigma}_l^2\mathbf{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  $\sigma_l = 0$  forces  $\omega_{l}$  to be zero, causing non-stationary components to be constant and weights to act stationary, resulting in strict equivariance. Treating lengthscales  $\omega$  as learnable hyperparameter has the additional advantage that it allows learning of convolutional filter frequencies from data through model selection of hyperparameter  $\bar{\sigma}_l^2$ .
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  $\sigma_{l} = 0$  limiting case and relaxed layer-wise equivariance at higher variances  $\sigma_l > 0$ . Unlike prior works [Finzi et al., 2021a, van der Ouderaa et al., 2022] that require setting or tuning the prior variances  $\sigma_l^2$  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  $\eta_{*} = \arg \max_{\eta} p(\mathcal{D}|\eta)$ , whilst integrating out the  $P$  model parameters  $\pmb{\theta} \in \mathbb{R}^{P}$ :
$$
p (\mathcal {D} | \boldsymbol {\eta}) = \int_ {\boldsymbol {\theta}} p (\mathcal {D} | \boldsymbol {\theta}, \boldsymbol {\eta}) p (\boldsymbol {\theta}) \mathrm {d} \boldsymbol {\theta}
$$
The marginal likelihood  $p(\mathcal{D}|\boldsymbol{\eta})$  forms the normalising constant of the unnormalised posterior over parameters:
$$
p (\pmb {\theta} | \mathcal {D}, \pmb {\eta}) = \frac {p (\mathcal {D} | \pmb {\theta} , \pmb {\eta}) p (\pmb {\theta})}{p (\mathcal {D} | \pmb {\eta})} = \frac {p (\mathcal {D} | \pmb {\theta} , \pmb {\eta}) p (\pmb {\theta})}{\int_ {\pmb {\theta}} p (\mathcal {D} | \pmb {\theta} , \pmb {\eta}) p (\pmb {\theta}) \mathrm {d} \pmb {\theta}}
$$
As the marginal likelihood is constant in  $\theta$ , it can be ignored when training neural network parameters with a regular MAP point-estimate  $\theta_{*} = \arg \min_{\theta}\mathcal{L}_{\theta}$ :
$$
\mathcal {L} _ {\boldsymbol {\theta}} = - \log p (\mathcal {D} | \boldsymbol {\theta}, \boldsymbol {\eta}) - \log p (\boldsymbol {\theta})
$$
To optimise hyperparameters, the marginal likelihood  $p(\mathcal{D}|\boldsymbol{\eta})$  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  $\pmb{\mu} = \pmb{\theta}_{*}$  and the covariance set to the local curvature  $\pmb{\Sigma} = \mathbf{H}_{\pmb{\theta}_{*}}^{-1} \coloneqq -\nabla_{\pmb{\theta}}^{2}\mathcal{L}_{\pmb{\theta}}|_{\pmb{\theta} = \pmb{\theta}_{*}}$  at that point:
$$
p (\boldsymbol {\theta} | \mathcal {D}, \boldsymbol {\eta}) \approx q (\boldsymbol {\theta} | \mathcal {D}, \boldsymbol {\eta}) = \mathcal {N} (\boldsymbol {\mu}, \boldsymbol {\Sigma})
$$
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  $\eta_{*} = \arg \min_{\eta}\mathcal{L}_{\eta}$
$$
\mathcal {L} _ {\boldsymbol {\eta}} = - \log \left(\sqrt {\frac {(2 \pi) ^ {P}}{| \boldsymbol {H} _ {\boldsymbol {\theta} _ {*}} |} p (\mathcal {D}, \boldsymbol {\theta} _ {*} | \boldsymbol {\eta})}\right) = \underbrace {- \log p (\mathcal {D} | \boldsymbol {\theta} _ {*} , \boldsymbol {\eta})} _ {\text {N L L / D a t a f i t}} - \underbrace {\log p (\boldsymbol {\theta} _ {*}) - \frac {P}{2} \log (2 \pi) + \frac {1}{2} \log | \boldsymbol {H} _ {\boldsymbol {\theta} _ {*}} |} _ {\text {O c c a m ' s f a c t o r}} \tag {9}
$$
Although the Laplace approximation greatly simplifies the marginal likelihood, the Hessian determinant  $|\pmb{H}_{\theta_*}|$  remains difficult to compute for large neural networks due to quadratic scaling in parameter count  $\mathcal{O}(|\pmb{\theta}|^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  $\mathcal{L}_{\theta}$  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  $\mathcal{L}_{\eta}$  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  $\mathcal{L}_{\eta}$  as objective to learn layer-wise equivariances from data. As Hessians scale quadratically in the number of parameters
$\mathcal{O}(|\theta|^2)$ , the loss becomes intractable for large networks. We use KFAC [Martens and Grosse, 2015] to approximate the Hessian of the log likelihood  $H_{\theta_*}$  and have the Gauss-Newton approximation  $\pmb{H}$ :
$$
\boldsymbol {H} _ {\boldsymbol {\theta} _ {*}} \approx \boldsymbol {H} = \sum_ {n} \boldsymbol {H} _ {n} = \sum_ {n} \boldsymbol {J} (\boldsymbol {x} _ {n}) ^ {\top} \boldsymbol {\Lambda} (\boldsymbol {x} _ {n}) \boldsymbol {J} (\boldsymbol {x} _ {n}),
$$
where the Jacobians of network outputs  $\pmb{f}$  with respect to parameters  $\theta$  for data points  $\pmb{x}_n$  are  $[J(\pmb{x}_n)]_{k,p} = \frac{\partial f_k}{\partial\theta_p}(\pmb{x}_n)$ , and  $[\Lambda(\pmb{x}_n)]_{k,g} = \frac{\partial\log p(y_n|\pmb{x}_n;\hat{\pmb{f}}(\pmb{x}_n;\theta))}{\partial f_k\partial f_g}$  is the log likelihood Hessian, all evaluated at  $\theta_*$ . KFAC is a block-diagonal approximation of the GGN, with individual blocks  $\pmb{H}_l = \sum_n\pmb{H}_{n,l} = \sum_n\pmb{J}_l(\pmb{x}_n)^\top\pmb{\Lambda}(\pmb{x}_n)\pmb{J}_l(\pmb{x}_n)$  for each layer  $l$  with parameters  $\theta_l$  and Jacobian  $\pmb{J}_l(\pmb{x})$ .
For a fully-connected layer,  $\pmb{\theta}\pmb{x}$  with  $\pmb{\theta} \in \mathbb{R}^{G_l \times D_l}$  and  $\pmb{x} \in \mathbb{R}^{D_l}$ , we can write Jacobians of a single data point w.r.t. parameters as  $\pmb{J}_l(\pmb{x}_n)^\top = \pmb{a}_{l,n} \otimes \pmb{g}_{l,n}$ , where  $\pmb{a}_{l,n} \in \mathbb{R}^{D_l \times 1}$  is the input to the layer and  $\pmb{g}_{l,n} \in \mathbb{R}^{G_l \times K}$  is the transposed Jacobian of  $\pmb{f}$  w.r.t. the  $l$ th layer output. We thus have:
$$
\begin{array}{l} \pmb {H} _ {l} = \sum_ {n} \pmb {J} _ {l} (\pmb {x} _ {n}) ^ {\top} \pmb {\Lambda} (\pmb {x} _ {n}) \pmb {J} _ {l} (\pmb {x} _ {n}) = \sum_ {n} [ \pmb {a} _ {l, n} \otimes \pmb {g} _ {l, n} ] \pmb {\Lambda} (\pmb {x} _ {n}) [ \pmb {a} _ {l, n} \otimes \pmb {g} _ {l, n} ] ^ {\top} \\ = \sum_ {n} [ \boldsymbol {a} _ {l, n} \boldsymbol {a} _ {l, n} ^ {\top} ] \otimes [ \boldsymbol {g} _ {l, n} \boldsymbol {\Lambda} (\boldsymbol {x} _ {n}) \boldsymbol {g} _ {l, n} ^ {\top} ] \approx \frac {1}{N} \left[ \sum_ {n} \boldsymbol {a} _ {l, n} \boldsymbol {a} _ {l, n} ^ {\top} \right] \otimes \left[ \sum_ {n} \boldsymbol {g} _ {l, n} \boldsymbol {\Lambda} (\boldsymbol {x} _ {n}) \boldsymbol {g} _ {l, n} ^ {\top} \right]. \\ \end{array}
$$
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,  $\pmb{J}_l(\pmb{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  $\theta_{1}$  and  $\theta_{2}$  as a matrix multiplication 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  $\theta_{2}$  and then  $\theta_{1}$ . Each of these operations can be written as a matrix product  $\theta x$  with corresponding appropriate sizes of  $\theta$  and  $x$ . In comparison to the standard KFAC case, we have a Jacobian that can be written as
$$
\boldsymbol {J} _ {\boldsymbol {\theta}} ^ {\top} = \sum_ {c} \boldsymbol {x} _ {c} \otimes \left[ \nabla_ {\mathbf {y} _ {c}} \boldsymbol {f} \right] ^ {\top} = \sum_ {c} \boldsymbol {a} _ {c} \otimes \boldsymbol {g} _ {c},
$$
which means an additional sum needs to be handled. In particular, plugging the Jacobian structure into the standard KFAC approximation, we have
$$
\begin{array}{l} \boldsymbol {H} _ {\boldsymbol {\theta}} \approx \sum_ {n} \boldsymbol {J} _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n}) ^ {\top} \boldsymbol {\Lambda} (\boldsymbol {x} _ {n}) \boldsymbol {J} _ {\boldsymbol {\theta}} (\boldsymbol {x} _ {n}) = \sum_ {n} \left[ \sum_ {c} \boldsymbol {a} _ {n, c} \otimes \boldsymbol {g} _ {n, c} \right] \boldsymbol {\Lambda} (\boldsymbol {x} _ {n}) \left[ \sum_ {c} \boldsymbol {a} _ {n, c} \otimes \boldsymbol {g} _ {n, c} \right] ^ {\top} \\ = \sum_ {n, c, c ^ {\prime}} [ \pmb {a} _ {n, c} \pmb {a} _ {n, c} ^ {\top} ] \otimes [ \pmb {g} _ {n, c ^ {\prime}} \pmb {\Lambda} (\pmb {x} _ {n}) \pmb {g} _ {n, c ^ {\prime}} ^ {\top} ] \approx \frac {1}{N C} [ \sum_ {n, c} \pmb {a} _ {n, c} \pmb {a} _ {n, c} ^ {\top} ] \otimes [ \sum_ {n, c} \pmb {g} _ {n, c} \pmb {\Lambda} (\pmb {x} _ {n}) \pmb {g} _ {n, c} ^ {\top} ], \\ \end{array}
$$
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  $\sigma_{l}$  values. In App. C.1, we derive the KFAC for  $\theta_{1}$  and  $\theta_{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  $\pmb{a}$  parameters, instead of the induced weights  $s^{\omega}$ , we apply the chain rule:
$$
\frac {\partial \boldsymbol {f}}{\partial a _ {j}} = \frac {\partial \boldsymbol {f}}{\partial \boldsymbol {s} ^ {\omega}} \frac {\partial \boldsymbol {s} ^ {\omega}}{\partial a _ {j}}, \text {u s i n g t h e p a r t i a l d e r i v a t i v e} \quad \frac {\partial \boldsymbol {s} ^ {\omega} \left(c ^ {\prime} , c , x , y\right)}{\partial a _ {j}} = \phi_ {j} ^ {c ^ {\prime}, c} (x, y)
$$
Consequently, we can compute KFAC for sparsified layers by projecting the appropriate dimensions of Kronecker factors of existing KFAC derivations with basis functions  $\phi$  stored in the forward pass. In App. C.2, we derive KFAC for sparsified S-FC and S-CONV layers in detail.
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.
Example:
(3, upper left)
Symmetry	Prediction task	MAP	Diff. Laplace
FC 112.7 M	CONV 0.4 M	FC 112.7 M	CONV 0.4 M	Learned (ours) 1.8 M
Strict symmetry	(digit)	95.73	99.38	97.04	99.23	98.86
Partial symmetry	(digit, quadrant)	95.10	24.68	95.46	24.50	99.00
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  $\mathrm{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  $\theta$ , anchor points  $\pmb{u}$ , or lengthscales  $\omega$ , as discussed in Section 3.2. As expected, CONV layers perform best under regular MAP training, followed by less constrained FC and  $\mathrm{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.
Table 1: Preventing symmetry misspecification. Test accuracy for non-symmetric FC, strictly symmetric CONV, and learnable symmetry F-FC+CONV models on a translation invariant task and a 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.
Layer	Learnable equivariance	Prior N(·|η)	# Params	MAP	Diff. Laplace (Ours)
Test NLL (↓)	Test Acc. (↑)	Test NLL (↓)	Test Acc. (↑)	Lη(↓)
FC	θ	133.6 M	2.423	64.42	1.319	52.36	1.896
CONV	θ	0.3 M	1.184	82.81	0.464	84.07	1.022
FC+CONV	✓	θ	133.9 M	1.713	76.93	0.489	83.32	1.019
F-FC	θ	1.5 M	4.291	50.47	1.307	53.38	2.119
F-FC+CONV (Ours)	✓	θ	1.8 M	1.343	81.61	0.468	83.92	1.277
S-FC	u	0.8 M	5.614	50.21	1.496	46.18	1.926
S-CONV	u	0.2 M	0.604	82.53	0.562	81.29	1.037
S-FC+S-CONV (Ours)	✓	u	1.0 M	5.485	58.80	0.538	81.55	0.978
S-FC	ω	0.8 M	7.783	49.69	1.319	52.22	1.813
S-CONV	ω	0.2 M	0.913	81.84	0.549	81.20	1.055
S-FC+S-CONV (Ours)	✓	ω	1.0 M	2.040	77.97	0.564	80.63	1.067
Table 2: Equivariant convolutional layers CONV obtain the highest performance on image classification 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.
Layers learn to become convolutional, except at the end of the network. Upon further inspection of learned symmetries, we observe very high prior precisions  $\sigma_l^{-2} > 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  $1 \times 1$  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 a per-layer basis.
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+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 ( $\alpha = 10$  to  $\alpha = 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.
Dataset	MAP	Learned with Differentiable Laplace (ours)	Rel. Effective Num. of Param.
# Params	Test NLL (↓)	Test accuracy (↑)	Test NLL (↓)	Test accuracy (↑)	Approx. MargLik (↓)	FC (%)	CONV (%)	GCONV (%)
MNIST	1.2 M	0.172	97.59	0.023	99.21	0.328	10	(0-46)	15	(0-98)
Translated MNIST	1.2 M	0.812	90.78	0.053	98.27	0.216	0	(0-0)	23	(0-99)
Rotated MNIST	1.2 M	0.819	91.02	0.136	95.55	0.896	8	(0-20)	8	(0-47)
CIFAR-10	1.2 M	3.540	68.33	0.552	80.94	0.926	0	(0-1)	44	(0-99)
Rotated CIFAR-10	1.2 M	5.953	48.30	1.236	55.68	1.630	4	(0-22)	14	(0-41)
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.
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.