Learning Invariant Weights in Neural Networks
Abstract
Assumptions about invariances or symmetries in data can significantly increase the predictive power of statistical models.
Many commonly used machine learning models are constraint to respect certain symmetries, such as translation equivariance in convolutional neural networks, and incorporating other symmetry types is actively being studied.
Yet, learning invariances from the data itself remains an open research problem.
It has been shown that the marginal likelihood offers a principled way to learn invariances in Gaussian Processes.
We propose a weight-space equivalent to this approach, by minimizing a lower bound on the marginal likelihood to learn invariances in neural networks, resulting in naturally higher performing models.
INTRODUCTION
Intuitively, invariances allow models to extrapolate, or rather ‘generalise’, beyond training data (see Figure 1 for an extreme example).
An invariant model does not change in output when the input is changed by transformations to which it is deemed invariant.
The most straightforward way to achieve this, is perhaps by enlarging the dataset with transformed examples: a process known as data augmentation.
A link between invariance and data augmentation in kernel space was made by Dao et al.
We show that this invariance can equivalently be described as transformations on the weights, similar to Cohen and Welling [2016] where a neural network is constrained to respect rotational symmetry through rotated weight copies.
We do what is common in Bayesian model selection and find the correct invariance using the marginal likelihood. Optimizing the marginal likelihood has proven an effective way to learn invariances in Gaussian Processes (GPs)
[van der Wilk et al., 2018], but is not tractable for commonly used neural networks.
To overcome this, we propose a lower bound of the marginal likelihood capable of learning invariances in neural networks.
By learning distributions on affine groups, we can select the correct invariance for a particular task, without having to perform cross-validation or even requiring a separate validation set.
We succesfully learn the correct invariance on different MNIST and CIFAR-10 image classification tasks leading to better performing models.
RELATED WORK
Convolutional neural networks (CNNs) have been successful in a wide range of problems and played a key role in the success of Deep Learning [LeCun et al., 2015].
It is commonly understood that the translational symmetries that arise from effective weight-sharing in CNNs is an important driver for its outstanding performance on many tasks.
In Cohen and Welling [2016], a group-theoretical framework was proposed extending CNNs beyond translational symmetries, and demonstrated this for discrete group actions.
Many studies since have proposed ways to incorporate other symmetries in neural network weights, such as continuous rotation, scale and translation, into the weights of neural networks [Worrall et al., 2017, Weiler et al., 2018, Marcos et al., 2017, Esteves et al., 2017, Weiler and Cesa, 2019, Bekkers, 2019] and recent efforts allow practical equivariance in neural networks for arbitrary symmetry groups [Finzi et al., 2021].
Nevertheless, weight symmetries are typically fixed, must be known in advance, and can not be adjusted.
Some studies have proposed invariance learning with data augmentations [Cubuk et al., 2018, Lorraine et al., 2020], but thus do not embed symmetry in weights and often require a validation loss.
[Zhou et al., 2020] do learn invariances as weight-sharing, but require a meta-learning procedure with an additional validation loss.
Benton et al.
[2020] circumvent the need for validation data by learning a distribution of input transformations directly on the training loss. But, in doing so rely on an additional explicit regularization
[Figure]
term that depends on how invariances are parameterised.
Similar to this work, Schwöbel et al.
[2020] propose to use a lower bound, but again only considers a distribution in the input space rather than on weights.
We learn invariant weights by optimizing the marginal likelihood: the common method in Bayesian statistics to perform model selection, which is parameterization independent with the aim of being generally applicable to any chosen parameterization of invariance.
Interestingly, it has been shown that the marginal likelihood objective coincides with an exhaustive leave-p-out cross-validation averaged over all values of p and held-out test sets [Fong and Holmes, 2020].
Lastly, in Topological VAEs [Keller and Welling, 2021] capsules with ‘rolling’ feature activations show similarities to the deterministically sampled features obtained from our method, but differ in the reliance on ‘temporal coherence’.
ON INVARIANT MODELLING
A model f ( · ) is deemed ‘strictly invariant’ its output is unaffected by a set of transformations: f ( T g ◦ x ) =
f
( x ) , ∀ g ∈ G, x ∈ X where each transformation T g is governed by a group action g ∈ G forming a group G .
We can obtain an invariant model by averaging model outputs over all transformations T g .
Although group theory introduces a rigid mathematical framework that is often used to describe and incorporate symmetries in statistical and machine learning models, it is restricted in the sense that the set of transformations that generate a group is always closed, by the definition of a group.
To illustrate, imagine the classic MNIST image recognition problem [LeCun et al., 1998]: here invariance to rotations up to a certain angle allows for better extrapolation to tilded versions of fitted digits and thus more robust predictions and increased sample efficiency.
However, invariance to full 360 degree rotations (all SO(2) group actions) may prohibit us from differentiating between a ‘6’ and a ‘9’.
In an effort to overcome this issue, we follow Dao et al.
Raj et al.
[2017], van der Wilk et al.
[2018], Benton et al.
[2020] and construct our invariant function f θ ( x ; η ) from a non-invariant function g θ ( x ) by summing over the orbit:
[Equation]
where p η ( T ) denotes a density over the group action transformations parameterised by a vector η .
Through this construction, we hope to induce a relaxed notion of invariance upon the model, sometimes referred to as ‘insensitivity’ [
van der Wilk et al., 2018], ‘soft-invariance’ [Benton et al., 2020], or ‘deformation stability’ [Bronstein et al., 2021].
The special case in which the density p η ( T ) is uniformly distributed over the orbit results in the ‘Reynolds operator’ from Group Theory, which averages functions and thereby induces a ‘strict invariance’ over the entire group.
3.1 INVARIANT SHALLOW NEURAL NETWORK
We construct our invariant function from a single-layer noninvariant neural network:
[Equation]
2) where σ ( · ) is the soft-argmax function, x is the input, and
W 1 and W 2 are the respective first and second layer weights and biases.
We omitted the bias terms for notational clarity.
In this study, we consider two flavours for our neural network g θ , namely an RFF-network and ReLU-network.
In the RFF-network, first layer weights W 1 are initialiased as Random Fourier Reatures (RFF)
[Rahimi et al., 2007] and a cosine activation function ϕ ( · ) =
cos( · ) is used.
For the ReLU-network, both first and second layer weights Wand W 2 are learned and we consider a ReLU non-linearity ϕ ( x ) = max(0 , x ) for the activation function.
The RFF-network is interesting because we obtain a weightspace equivalent that is as close as possible to a GP with a radial basis function kernel (RBF), with exact correspondence in the infinite-width limit.
From van der
Wilk et al.
[2018], we know that in this case the marginal likelihood is tight and can be used to learn invariance.
The ReLUnetwork, on the other hand, is interesting as it more closely resembles the commonly used architectures in the Deep Learning (DL) community: basis weights are typically not fixed and the ReLU is one of the most commonly used activation functions in DL.
In our experiments, we find that we can learn invariances with both the RFF-network and ReLU-network, indicating that for our purposes the bound on the marginal likelihood remains sufficiently tight for more general shallow architectures. ± ± ± ± ± ±
Section 3.7 will discuss how variational inference is used to learn a variational distribution q over the parameters in the second layer: θ = vec ( W 2 ) .
3.2 INVARIANCE IN THE WEIGHTS
In Equations 1-2, we showed how we construct an invariant function by integrating or summing over transformed input samples T ( x ) .
Yet, instead of explicitly performing these transformations on the input, we can obtain a mathematically equivalent invariant function by considering transformations on the weights.
Note that the inner term of our neural network definition in Equation 2, we have that ( W 1 ◦ T ) ◦ x =
W 1 ◦ ( T ◦ x ) are equal, by associativity of matrix transformations.
In other words, first applying transformation T on the weights, similar to the typical construction of equivariant layers, is equivalent to first applying it to the input, which could be interpret as built-in data augmentation.
In practice, however, differences between the two could still arise if applying T requires approximations (e.g. interpolation between discrete pixels).
In our experiments we will consider transforming the weights, thus demonstrating that invariance can be ‘built into’ the model.
Coordinate data and imaging data We consider simple affine transformations, which can be represented as T ∈ R 3 × 3 matrices.
For 2d vector data, applying the transformations amounts to regular matrix multiplications, which only requires appending a single 1 -entry to the data vectors.
For 2d images, where data points x ∈ R WH correspond to
W × H pixel grids, applying T in image space requires interpolation.
Here, we could use bilinear interpolation, which can also be written in matrix formula form.
We use the grid sample operation (as used in [Jaderberg et al., 2015]) which acts on the weight matrix values W 1 and outputs an equally shaped matrix. The operation treats the HW -dimensional row vectors of W 1 as a grid of H × W points where the coordinates of the point are transformed according to the affine transformation matrix T . The resulting values are obtained by interpolating the values of transformed pixels at the original grid coordinates using bilinear interpolation. 3.3 AFFINE LIE GROUP REPARAMETERIZATION
The transformations that are applied on the weights and will define the invariances of the network are sampled from a probability distribution. To allow learnable invariances, we define a learnable probability distribution over the transformations p η ( T ) parameterised by η .
We will refer to η as the ‘invariance parameters’, as they parameterise to which transformations to which our network becomes invariant.
To learn this distribution with back-propagation, we must make sure that samples taken from the distribution are differentiable with respect to the invariance parameter η .
For affine transformed weights, we consider a procedure similar to what Benton et al.
[2020] used to augment inputs, utilising the re-parameterization trick [Kingma and Welling, 2013] to remain differentiable. The distribution defines independent Gaussian probabilities over infinitesimal generators around their origin. sampling noise from a k-cubed uniform distribution ϵ ∼ U [ − 1 , 1] k .
With k = 6 generator matrices G 1 , · · · , G 6 and learnable parameters η
= [ η 1 , · · · , η 6 ] T we can separately parameterise translation in x, translations in y, rotations, scaling in x, scaling in y, and shearing (see Appendix D).
A sample T ∼ p η ( T ) can be obtained by transforming noise ϵ :
[Equation]
with matrix exponential exp( M ) = (cid:80) ∞ n =0 1 n !
M n .
A distribution over the subgroup of 2d rotations SO(2) can be achieved by only learning the parameter for rotational invariance η rot =
η 3 and fixing η i =
for all i ̸ = 3 .
Then,
[Equation]
[Equation]
By learning η rot , we can effectively interpolate between no invariance at η rot = 0 to full rotational invariance at η rot ≡ π .
E p o c h 1 Initial Invariance: ± 5°
Base filter
E p o c h 5 Learned Invariance: ± 81°
E p o c h 7 Learned Invariance: ± 115°
F i n a l E p o c
h Learned Invariance: ± 179° (a) Feature bank #1 over training iterations.
E p o c h 1 Initial Invariance: ± 5°
Base filter
E p o c h 5 Learned Invariance: ± 81°
E p o c h 7 Learned Invariance: ± 115°
F i n a l E p o c
h Learned Invariance: ± 179° (b) Feature bank #2 over training iterations.
Similarly, we can define a distribution over the subgroup of 2d translations T (2) by fixing η i = 0 for all i > 2 and learning the translational invariance parameters η 1 and η 2 :
[Equation]
We include full derivations including scaling in Appendix E.
In general, η = 0 corresponds to no invariance and increasing individual elements of η also increases insensitivity to corresponding transformations towards full invariance.
3.4 STOCHASTIC OR DETERMINISTIC SAMPLING
To estimate f θ ( x ; η ) from Equation 1, we approximate the integral with a Monte Carlo (MC) estimate:
[Equation]
where S transformations are stochastically sampled from the distribution T i ∼ p η ( T ) . Samples can be differentiated with respect to invariance parameter η using the ‘reparameterization trick’ (see Section 3.3).
We know that MC is an unbiased estimator, and thus
[Equation]
with E T :
= E (cid:81) Si =1 p η ( T i ) .
Unlike stochastic MC sampling, we can obtain a deterministic surrogate of the procedure by replacing the stochastic samples from the noise source U [ − 1 , 1] k with linearly spaced points along its k -cubed domain.
This procedure is similar to quadrature in classical numerical integration, or from a programming perspective, as applying the re-parameterization trick on a fixed ‘linspace’ instead of uniform noise.
A visualization of a discretely sampled filter bank of a model learning rotational invariance over training iterations is shown in Figure 3.
By ensuring sufficient and equally spaced samples, deterministic sampling can be used to ensure reliable and robust inference at test time.
Similar to the stochastic sampling, this deterministic procedure is also differentiable and can thus be used during training. We find, however, that deterministic sampling is only suitable when the number of invariances dim( η ) is very small (see Section 3.5).
Nevertheless, deterministic sampling can be theoretically interesting and allow our model to be interpret as a generalization of other architectures.
For instance, a single convolutional layer where the kernel is discretely and deterministically convolved over an image followed by spatial pooling, can be interpret as an instance of our invariant MLP with a specific affine invariance transformation in which weights are ‘zoomed-in’ and deterministically sampled and reapplied over the image plane.
3.5 PRACTICAL TRANSFORMATION SAMPLING
If η comprises multiple elements, the sampling suffers from the curse of dimensionality as the number of required samples grows exponentially with larger K .
To illustrate, a sparse 3 quadrature points in K = 6 dimensions would already require 3 6 = 729 samples with deterministic sampling.
In general, we found that stochastic MC sampling resulted in the most stable training behaviour and therefore used this when training the models in the experimental section, except for Figure 3 where deterministic sampling was used for both training and visualization of rotationally invariant filter bank.
[Figure]
3.6 LOWER BOUNDING THE MARGINAL LIKELIHOOD
We have a (typically large) vector θ containing the model parameters and a (typically small) vector for the invariance parameters η .
The approach we take in this paper is to perform Bayesian Model Selection and integrate out θ but find a point-estimate over η :
[Equation]
where p ( D| η ) is the marginal likelihood [Murphy, 2012] or model evidence, sometimes called empirical Bayes or type-II ML.
The technique has been shown effective in GPs to learn hyper-parameters Williams and Rasmussen [2006] and invariances van der Wilk et al.
[Equation] but is typically intractable for neural networks.
We derive a lower bound that allows for optimization of the marginal likelihood in neural networks using stochastic variational inference:
[Equation]
E T (cid:104) log p ( y | ˆ f θ ( x ; η ))
(cid:105)(cid:105) − KL
[Equation]
[Equation] [Equation] [Equation] 9) with expectations E θ :=
E q ( θ ) and E T :=
E (cid:81) Si =1 p η ( T i ) .
We begin Eq 9 with the standard evidence lower bound (ELBO) derived from variational inference.
In the second and third line, we expand the likelihood and plug-in Equation 7.
In the last line, we use Jensen’s inequality together with the fact that our log-likelihood is a convex function.
The resulting lower bound comprises an expected log-likelihood term that can be estimated by taking the average cross-entropy on mini-batches of data (see Section 3.7) and a KL-divergence between two multivariate Gaussians which can efficiently be computed in closed-form.
Note, we integrate out model parameter vector θ , which is part of the KL-term, whereas the vector parameterizing the invariances η is only part of the first term.
We optimise the derived lower bound w.r.t.
both η and θ every iteration with stochastic gradient descent.
3.7 VARIATIONAL INFERENCE
To summarise, we propose to learn invariances using stochastic variational inference [Hoffman et al., 2013] and derived a lower bound of the marginal likelihood, or evidence lower bound (ELBO) that can be optimised using a gradient descent methods, such as Adam [Kingma and Ba, 2014].
Variational inference minimises the KL-divergence between an variational posterior and the true posterior on our free model parameters p ( θ |D ) , where θ = vec ( W 2 ) .
For the approximate posterior, we choose a multivariate Gaussian distribution q ( θ | µ , Σ ) : = N ( θ | µ , Σ ) parameterised by variational parameters µ and block-diagonal covariance Σ with a separate block for each output class.
The covariance is parameterised as a Cholesky decomposition Σ
= L T L , which is a common trick to maintain computational stability to ensure a positive semi-definite Σ and does not influence the model.
We obtain a differentiable Monte Carlo estimate of q ( θ ) by sampling L times from the variational distribution, using the reparameterization trick [Kingma and Welling, 2013], and maximise the ELBO:
[Equation]
where we can choose L { = } 1 given a sufficiently large batch size.
We obtain a Stochastic Gradient Variational Bayes (SGVB) estimate of the lower bound
NM (cid:80) Mi =
[Equation]
˜ L ( θ , { x i } , { y i } )
[Kingma and Welling, 2013] to allow efficient training on mini batches of data.
Full derivations can be found in Appendix A.
[Table]
EXPERIMENTS AND RESULTS
We implemented our method in PyTorch [Paszke et al., 2017] and show results on a toy problem with different degrees of rotational invariance in Figure 1 with 1024 RFF features, σ = 5 , and T applied on the weights.
The following sections will describe experiments on different MNIST and CIFAR-10 image classification tasks where
T is applied on the weights by using the bilinear grid resampling as described in Section 3.2 and Jaderberg et al.
[2015] in combination with small 0.1 sigma Gaussian blur to bandlimit high frequencies.
We used Adam [Kingma and Ba, 2014] for optimization in combination with a learning rate of 0.001 ( β 1 = 0 .
9 , β 2 = 0 .
cosine annealed
[Loshchilov and Hutter, 2016] to zero.
Parameters were initialised as µ c = 0 , L c =
I for all classes c , σ = 0 . 3 , and α = 1 . 0 .
We use S = 32 samples from p η ( T ) , L = 1 and a batch size of 128.
4.1 ON THE NECESSITY OF A BAYESIAN APPROACH
To investigate to what extent the variational inference is required to learn invariances, we compare our approach with regular maximum likelihood using Adam. We train one model that uses our objective (Variational Inference), and another model where we replaced the variational distribution q ( θ ) with a point-estimate and omitted the KLterm to get a regular cross-entropy loss.
Interestingly, when trained on fully-rotated MNIST in Figure 5, we find that the model trained with cross-entropy was completely incapable of learning the correct invariance, whereas our VI-based approach does learn the invariance.
We hypothesise that maximum likelihood alone is not enough to learn invariance, as invariance is a constraint on the weights and thus does not help to fit the data better, whereas marginal likelihood also favours simpler models.
This result substantiates the use of marginal likelihood (or a lower bound thereof) for hyper-parameter selection for neural networks, and invariance learning in particular.
More broadly speaking, it proves a convincing case for probabilistic machine learning models, such as Bayesian neural networks, beyond their oft-cited use for uncertainty estimation.
[Figure]
4.2 IDENTIFYING INVARIANCE WITH ELBO
To evaluate whether the ELBO is capable of identifying the apt level of invariance, we consider models with different fixed values of rotational invariance η rot and one model where η rot is learned.
We then evaluate the models on three different versions of MNIST on which we artifically imposed different amounts of rotational invariance by randomly transforming the dataset beforehand. In ‘fully rotated MNIST’, we rotate every image with a random uniformly sampled angle in range [ − 180 ◦ , 180 ◦ ] . In ‘partially rotated
MNIST’ images are rotated with a random angle within [ − 90 ◦ , 90 ◦ ] . Lastly, we also consider the ‘regular MNIST’ dataset without any alterations.
From Table 1, we observe that for each dataset the model with the best ELBO corresponds to the model with the right amount of invariance, also correlating with best test accuracy.
This finding indicates that the ELBO can correctly identify the required level of invariance, and confirms that choosing the right invariance leads to better generalization on the test set.
On regular MNIST, we observe that a small amount of invariance yields better ELBO than no invariance. This could be explained by some intrinsic rotational variation within the dataset.
Furthermore, we find that the ELBO of the model with learned invariance η rot corresponds to the optimal ELBO in the set of models with fixed invariance.
Therefore, we find that in this case, we can use the ELBO to learn invariances in a differentiable manner.
Additional results can be found in Appendix C.
4.3 RECOVERING INVARIANCE FROM INITIAL CONDITIONS
To investigate robustness to different initial conditions, we repeat the experiment where we learn invariance parameters η during during training on fully-rotated, partially rotated and regular MNIST data but with different initial values, corresponding to rotational invariance of [ ± 5 ◦ , ± 45 ◦ , ± 90 ◦ , ± 135 ◦ , ± 175 ◦ ] degrees.
Results of this experiments for the
RFF neural network are shown in Figure 2, and a similar figure for the ReLU neural network is attached in Appendix
C.2.
For most initial conditions, we observe that we can succesfully learn and recover the ‘correct’ amount of invariance for each dataset.
One exception being initial 175 ◦ degrees on partially rotated dataset, which suggests that training with low initial invariance could be advantageous in practice, for this method.
Nevertheless, we conclude that our model can recover invariance relatively robustly independent of initial conditions.
4.4 LEARNING INVARIANCE IN RELU NETWORK
So far, we have only considered the set-up where we learn the output layer W 2 and keep the first layer W 1 initialised as fixed RFF-features in combination with a cosine cos( · ) activation function.
We chose this fixed basis function model to ensure a sufficiently tight bound on marginal likelihood where the only source of looseness is the non-Gaussian likelihood.
Now, we will let loose of these constraints and consider a general single hidden layer neural network with ReLU non-linearity ϕ
( x ) = max( x, 0) with Xavier [Kumar, 2017] initialised weights and 1024 hidden units, where we learn both the input layer W 1 the output layer W 2 .
We optimise the model using the same variational inference procedure.
We find that we are still able to learn invariances in the setting where parameters of both input and output layer are learned (full comparison in Appendix C).
In Figure 3, we plot an illustration of a feature bank (row vector in W 1 with 7 samples equally spaced between − η rot and η rot and plotted over training iterations).
The top of the figure shows the randomly initialised features without any rotational invariance at the beginning of training.
After training on a fully rotated MNIST, the features converge to a particular filter with practically full ±179 ◦ rotational invariance, as shown on the bottom of the same figure.
4.5 OTHER TRANSFORMATIONS
To explore invariance to transformations other than rotation, we allow for different kinds of affine invariance transformations, namely rotation, translation, scale and full affine transformations (see Section 3.3). Again, we use the ReLUnetwork where both layers are learned.
In Table 2, we evaluate and compare models that can learn affine invariances with two non-invariant baselines, namely a regular Gaussian Process regression with RBF kernel baseline (SGPR) and a regular shallow neural network baseline (MLP).
We use SGPR as a reference, because we know the training procedure is reliable and to ensure enough capacity is given to the single layer MLP.
We separately trained the models on fully-rotated, translated, scaled and original versions of MNIST (see Appendix D for details).
We find that models with learned invariances (bottom four rows) outperform the model with no invariance (top two rows) in all cases.
As expected, a translationally invariant model performs better on a dataset that contains randomly translated examples, and similarly, the rotationally and scale invariant models perform best on the respective rotated and scaled versions of MNIST.
In line with our expectations, the model capable of learning affine invariances performs best overall.
Moreover, by inspecting the learned coefficients of η after training we verified that the learned transformations correspond to the dataset the it was trained on.
This can also be observed in Figure 4 by inspecting the resulting learned filter banks samples after training on different datasets.
[Table]
We repeated the same experiment on the CIFAR-10 dataset Krizhevsky et al.
[2009] and trained on fully-rotated, translated, scaled version and the original version of the
[Figure] CIFAR-10 dataset and plot test accuracies in Table 3.
We consistently find that the best performing models are those that are parameterised such that it can learn the invariance that corresponds the dataset, typically resulting in several percentage points of improved accuracy compared to the MLP baseline.
Furthermore, if we parameterise the MLP with the more general affine invariance, capable of expressing rotation, translation and scale invariances, we always achieve similar or improved results compared to model from the models parameterised with a a single invariance.
Similar to the MNIST experiments, we find that the MLP with general affine invariances can select the correct invariance based on the used training data.
Here we also verified this by inspecting the θ , and found that the learned invariance always matches the invariances that we expect for the corresponding dataset.
For example, the model capable of learning affine invariances correctly learned to be rotationally invariant ( η 3 ≈ π and η i ≈ 0 for i ̸
= 3 ) after training on the fully-rotated CIFAR-10 dataset.
[Table]
DISCUSSION AND CONCLUSION
In this paper, we propose a single training procedure capable that can learn invariant weights in neural networks automatically from data.
We follow what is common in Bayesian statistics and optimise the marginal likelihood to perform Bayesian model selection: a method that has been proven capable to learn invariances in GPs [van der
Wilk et al., 2018].
We propose a lower bound to allow optimization of the marginal likelihood in shallow neural networks. On MNIST and CIFAR-10 image classification tasks, we demonstrate that we can automatically learn weights that are invariant to correct correct affine transformations, solely using training data.
Furthermore, we show that this leads to better generalization and higher predictive test accuracies.
The marginal likelihood is a general model selection method and is parameterization independent.
Therefore, we can expect it to work on other invariances and other model architectures.
In this work, we focussed on affine transformations, but it would be interesting to consider more complex parameterizable transformations over the image space, such as diffeomorphic vector fields [Schwöbel et al., 2020].
We showed that we can learn invariance by sampling a learned compactly supported continuous probability distribution over group actions in common Lie groups.
Allowing discrete groups would either require differentiating through a discrete probability distribution, for instance utilizing the
Gumbel-Softmax trick [Jang et al., 2016]), or, by treating the discrete group as a subgroup of some Lie group and learn to approximately distribute all continuous density p η ( T ) to the group actions of the subgroup.
Furthermore, it would be interesting to consider more flexible and complex probability densities over group actions, such as a mixture distributions or normalizing flows [Rezende and Mohamed, 2015, Tabak and Turner, 2013], capable of expressing multiple modes as in [Falorsi et al., 2019].
We found that we could succesfully learn invariance using marginal likelihood, also referred to as Empirical Bayes or Type-II ML, which is not possible with regular maximum likelihood (Type-I ML).
To do so, we relied on η being small and learning higher dimensional invariances might therefore require more sophisticated methods or additional priors on η .
Lastly, this work focuses on single layer neural networks, and we will consider deeper architectures in future work.
For deeper models, we should ask the question whether the bound on the marginal likelihood will stay sufficiently tight [Dutordoir et al., 2021, Ober and Aitchison, 2020, Immer et al., 2021].
To conclude, we hope our findings inspire other works to allow neural networks that automatically learn symmetries from data.
[-t. parameters θ = vec ( W 2 ) by minimizing the D KL ( ·||· ) -divergence between approximate posterior q ( W 
-]
[-Fix η i = 
-]
[-Fix η i = 
-]
[-Fix η i = 
-]