Invariance Learning in Deep Neural Networks with Differentiable Laplace Approximations
Authors: Alexander Immer
Tycho F.A. van der Ouderaa
Gunnar Ratsch
Vincent Fortuin
Mark van der Wilk
Venue: NeurIPS 2022 Accept (2022)
Keywords: Bayesian deep learning
Laplace approximation
invariance learning
Bayesian model selection

Abstract: Data augmentation is commonly applied to improve performance of deep learning by enforcing the knowledge that certain transformations on the input preserve the output. Currently, the data augmentation parameters are chosen by human effort and costly cross-validation, which makes it cumbersome to apply to new datasets. We develop a convenient gradient-based method for selecting the data augmentation without validation data during training of a deep neural network. Our approach relies on phrasing data augmentation as an invariance in the prior distribution on the functions of a neural network, which allows us to learn it using Bayesian model selection. This has been shown to work in Gaussian processes, but not yet for deep neural networks. We propose a differentiable Kronecker-factored Laplace approximation to the marginal likelihood as our objective, which can be optimised without human supervision or validation data. We show that our method can successfully recover invariances present in the data, and that this improves generalisation and data efficiency on image datasets.

Invariance Learning in Deep Neural Networks with Differentiable Laplace Approximations
Alexander Immer Tycho F.A. van der Ouderaa +.3 Mark van der Wilk
Gunnar Rätsch
Vincent Fortuin
Department of Computer Science, ETH Zurich, Switzerland
2 Max Planck Institute for Intelligent Systems. Tübingen, Germany *Department of Computing, Imperial College London, UK Department of Engineering, University of Cambridge, UK
Abstract
Data augmentation is commonly applied to improve performance of deep learning by enforcing the knowledge that certain transformations on the input preserve the output. Currently, the data augmentation parameters are chosen by human effort and costly cross-validation, which makes it cumbersome apply to new datasets. We develop convenient gradient-based method for selecting the data augmentation without validation data during training of deep neural network. Our approach relies on phrasing data augmentation as an invariance in the prior distribution on the functions of neural network, which allows to learn it using Bayesian model selection. This has been shown to work in Gaussian processes, but not yet for deep neural networks. We propose a differentiable Laplace approximation to the marginal likelihood as our objective, which can be optimised without human supervision or validation data. We show that our method can successfully recover invariances present in the data, and that this improves gencralisation and data efficiency on image datasets.
cronecker-factoree
1 Introduction
Data augmentation is commonly used machine learning technique that is essential high- performing deep learning and computer vision systems. It aims to obtain model that is invariant to set or distribution of transformations, by fitting a model with inputs that are transformed in way that is known to leave the output class unchanged. This procedure can be regarded as artificially creating more data and is well known to increase generalisation performance and data efficiency. Yet, choosing the right transformations is an expensive and task-specific process that relies on domain knowledge and human effort, as well as trial-and-error through cross-validation. This can quickly become intractable when many parameters each setting requires training a model to convergence.
considered, particularly if they continuous, because
We aim to make selecting suitable transformations easier, by learning them via gradient descent. Our approach is inspired by the procedure of yaan der Wilk et al.) (2018), which casts learning invariances and data augmentations as a Bayesian model selection problem. This view suggests selecting invariances by maximising the marginal likelihood with gradient- based optimisation. While this approach was successful in Gaussian process models, it has not yet been demonstrated in deep neural networks, where the marginal likelihood is harder to approximate.
Equal contribution, order decided by coin flip. Correspondence alexander. immerdinf ethz. ch.
tycho. vanderouderaa@imprria uc
The code is available at phattps://giteee com/tychovdoollla
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

x
f(x, 0,), =
Eptxc(xn)) [((xx0)
-log
f(x, R.
100
DNN f
fsJ
Avg.
H50" learned invarisacee 201
p(x (x.
+80
Suyg 500
Step A: sample S augmented data points/scts dependent on invariance parameter ! (here controls max rotation)
Step B: obtain invariant deep
Step C: Approximate and diffcrentiate marginal likelihood wrt. 7) to optimize
ncural network by averaging over augmentation samples x' from p(x' x,
invariance during training
Figure 1:
Illustration of our approach on dataset with rotational invariance. Steps A and allow to define Bayesian neural network with a likelihood dependant on invariance parameter v/ (van der Wilk al. 2018). Our contributions enable tractable marginal likelihood estimation and differentiation during training, illustrated in Step C. The right-most figure shows the posterior predictive of our Bayesian neural network after learning the invariance present in the data.
To circumvent this problem, we built upon the scalable Laplace approximation to the marginal likelihood developed bylmmmrr et al. (2021a), who recently showed that maximising it can successfully select neural network hyperparameters, such as architectures. We extend their method to enable gradient-based optimisation of complex hyperparameters that control invariances in deep neural networks. To that end, we propose an efficient and differentiable Kronecker-facton facttord Laplace approximation for invariant neural networks and novel method obtain stochastic gradients with respect to invariance parameters, which is also useful for optimising other hyperparameters. Our method the first enable differentiable Bayesian model selection learn complex hyperparameters, in particular invariances, in deep neural networks.
Our approach is illustrated in Fig. We specify invariances parameterised distribution perturbations of the network's 's input, like in data augmentation (Step A). The model output is averaged over samples from the distribution (Step B), which yields a Bayesian neural network with likelihood dependent on the perturbations (Nabarro et al. 2021). We derive a marginal likelihood approximation for such neural networks and show how to efficiently compute its gradients with respect to the invariance parameters (Step C). By approximating the marginal likelihood, we can differentiably learn invariances during training, jointly with neural network parameters, and without validation data.
We demonstrate experimentally that our method can differentiably learn useful distributions over affine invariances, which are common data augmentations, datasets MNIST, FashionMNIST, and CIFAR-10, 10, without validation data. Our learned invariances improve the generalisation and data efficiency of neural networks, without the effort required for choosing data augmentations or custom architectures. On original datasets, our method can increase test performance by up 8-10 percentage points. On random subsets of image classification datasets, we show that our method can achieve up to 10x better data efficiency. Our work strengthens how Bayesian methods can be useful for deep learning beyond predictive uncertainty estimation.
various versions of the image classification
2 Related Work
Invariances in deep learning. Since the inception of the convolutional neural network (CNN)
(Fukushima and Miyake, T982, LeCun et al. 1998). building invariances and equivariances into
deep learning models has drastically increased data efficiency and generalisation (Cohen et al. 2018
Brandstetter et al. 202T). for instance, on image classification (Cohen and Welling 2016). molecular dynamics (Batzner et al., 2021), and reinforcement learning (van der Pol et al. 2020). However, these approaches require knowing the invariances a priori. this work, on the other hand, we aim to automatically learn the correct type and amount of invariance from data without supervision.
Learning invariances. Learning invariances from data is hard, because symmetries define con- straints on the functions a network can represent, and therefore do not improve data fit according to the training loss, even if they would lead to better generalisation have therefore proposed learn invariances or data augmentation by estimating gradients on validation loss. For example, AutoAugment (Cubuk et al. 2018) Iearns data augmentation using
test data. Some methods

---

policy gradients, Lorraine et al. (2020) use the implicit function theorem, and Zhou et al. (2020) phrase the problem as meta-learning. Such approaches require validation sets of sufficient size to prevent high variance and overfitting (Lorrinee et al. 2020)). Here, we tackle the problem of learning
invariances when validation data is not available.
Validation-free invariance learning. When casting invariance learning a Bayesian model selection problem, we can select helpful invariances using the marginal likelihood of the model with training data alone. This has been successfully demonstrated in Gaussian processes (van der Wilk et al., 2018) and in the weight space of single-layer neural networks (van der Ouderaa and van der Wilk 2021). To scale this approach to deep networks, Schwöbel et al. (2022) attempt to use marginal likelihood that is only computed in the last layer. While the latter was successful for small neural networks, it failed learn invariances on more complex datasets that require deep neural networks, likely due to known limitations of last-layer approaches (Ober et al. [021] van Amersfoort et al., (02T). Augerino (Benton al. 2020) is an alternative to Bayesian model selection that works for deep learning by regularising invariances to increase during training. However, this approach depends on the chosen parameterisation of the invariance and may still require a validation set to tune the regularisation strength. In App. C we show that these issues can be detrimental to Augerino's 's performance.
contrast, our proposed method uses the Bayesian marginal likelihood estimate in deep neural networks that is iaaannttriiationinnaaaaaa and improves performance. The main differences between our proposed method and the alternatives are summarised in Appp. A
Bayesian model selection for deep learning. Marginal likelihood approximations have recently en- abled gradient- based hyperparameter optimisation for deep neural networks (Immer et al. 2021a: la Ober and Aitchison 1 Antorán et al. 2022a). Laplace approximations (MacKay 1992) with structured Hessian approximations, such as KFAC (Martens and Grosse 2015 Botev et al. 2017). or linearisation have been shown to improve generalisation, for example, by optimising weight regularisation (Immer ctt al. 202 a Daxberger et al. (021) I) and learning length-scall scales of convolutions (Antorán et al., (202a)).
Other approximate inference methods that rely ensembling or sampling do not directly apply to
marginal likelihood estimation while common methods like mean-fiele field variational inference (Blundell
et al. have been shown to fail for model selection. We therefore focus on Laplace approximations and extend them to enable gradient-basse optimisation of more complex hyperparameters like invariances.
2015) last-layer approaches (Ober et al., [202 Daxberger et al. 2021 Schwööel et al. 2022)
3 Background
We consider supervised learning task with dataset D = (Xx, ((xx,..)... | consisting of N inputs x E TR. and targets y E R,C Our goal is learn the function f RR. - RC that relates the inputs and outputs, which we represent as neural network fxx; 9) with parameters O. We control which solutions for f are preferred over others (i.e., the inductive bias) with hyperparameters H. which parameterise the prior over 0 and f. Invariance is particularly helpful inductive bias that constrains the output of to remain similar for certain transformations of the input x, which can improve generalisation by allowing single datapoint to inform predictions for wider range of inputs. Our goal is to learn useful invariances together with the network weights.
3.1 Parameterising Invariance
To construct and parameterise invariant functions, we consider a local invariance that intuitively requires the function to not change "too much" in response to transformed inputs. To obtain such an invariant function f, we average an unconstrained function f over perturbation distribution p(x'|x [x, .7). In practice, we sample x' by reparameterising ~ p(c) with a differentiable function and approximate the expectation with S Monte Carlo samples , Hi p(e) (van der Wilk et al. 2018 Benton et al. 2020):
fxx; 0 n) Peptte (x.) [f(x': 0)] (P e; 0)) ~ E (g((x (.nk,,sm)))) m): 0),
where f is differentiable in the parameters 17 that control the perturbation distribution and therefore the invariance. When the perturbation distribution is uniform on the orbit of a group, we recover exact invariance in f (Kondor, 2008 Ginsbourger et al. (0166). The perturbation distribution resembles data augmentation applied to f instead of the loss and is also used at test time (f depends on it). Because the unconstrained function f is neural network, we refer to f an invariant neural network.
3

---

Given that we can parameterise suitable data augmentation distribution p(x' X, m),, we can learn invariances in arbitrary domains or group structures. What parameterisation works best in practice remains a research question. In our experiments, we consider a combination of uniform distributions and corresponding parameters 7 E R over generator matrices that define a probability density over the group of affine transformations, similar to Benton et al. (020) and detailed in App. B
Finding invariance parameters 17 is hard because f is constrained version of f and, especially for flexible models, this cannot improve the data fit according to the standard training loss (see also App. G.I. To overcome this, we propose to use Bayesian inference which provides convenient framework to optimise invariance parameters with gradients during training without validation data.
3.2 Bayesian Model Selection
Bayesian inference prescribes how unknowns, such as invariance parameters, should be determined from data. To infer hyperparameters H. from data, we interested in the their posterior, p(H)D) x p ((H (MacKay 2003). Because this posterior is intractable, type II maximum likelihood (MI.III) is often used instead, which obtains a point estimate 7t for the optimal hyperparameters according to the marginal likelihood ((D(H), which requires integration over the model parameters.
ML-II is routinely used in Gaussian processes [Rasmussen and Williams 2006, 5.2) and trades off model simplicity with data fit (Rasmussen and Ghahramani] (2001). There are also strong relations to quantities from other statistical frameworks, like cross- validatiin (Fong and Holmes (0200). Minimum Description Length (Grünwald 2007). and generalisation error bounds (Germain et al. 2016). Van der Wilk et al. (2018) showed its usefulness for selecting invariances, and elaborated on the mechanism by which it works. The main advantage of the marginal likelihood,
HDD = f p(D 0, H) H) do S p(yn (x,,; 0), H))((0)H)dO. do,
(2)
that it can be computed from the training data alone and optimised with gradients. In our work, the integration is over the neural network parameters and requires particularly scalable approximations.
3.3 Bayesian Model Selection for Deep Learning
Computing the marginal likelihood for neural networks involves intractable integrals. The Laplace approximation (Laplace 1774; MacKay 2003 $ 27) offers a solution by approximating the log joint of the parameters and data, ((2, 0 HH), with second-order Taylor expansion around mode O.:
log p (D)H) ss ((D, 10,,0,170 0, - 2 log L 2n |
where the first term decomposes into the log likelihood, a sum over data points, and the log prior, both of which are cheap and easy to evaluate and correspond to typical training loss evaluated at mode 9.. The second term depends on the log determinant of the log-joint Hessian at the same mode
Ho_ V2 p(D, 0 lo = H. v log (()))) (o-))loo...
(4)
where HNLL denotes the Hessian of the negative log likelihood. This approach allows to estimate the marginal likelihood using a MAP estimate 0, of the weights and its local curvature He..
To circumvent the high cost of estimating the full Hessian, structured generalised Gauss-Newtor (GGN) approximations are preferred for model selection in deep learning, as also in optimisa- tion (Martens 2020 BBttoo (010). [mmer et al. (2021a) recently demonstrated successful hyper- parameter and architecture selection with such approximations. Further, they observe empirically that their algorithm does not require to be at mode 0., which allows for interleaved gradient- based optimisation of parameters and hyperparameters during training. With the Jacobian matrix (x) det 0f(x:0)
80 E RCxP of the network output given input x w.r.t. parameters, and Hessian of the log likelihood w.r.t. network outputs Aff - V fogp(yyf] (f), the GGN simplifies the negative-log- likelibooo-ddependent term of the Hessian:
H R HoON
N Jo T O(((xx;)))Joo(x.)
(5)
Here, we assume that f forms the natural parameters of an exponential family likelihood func- tion (Murphy) 2012). In classification, for example, f are the logits. We refer to the resulting approximations as Laplace-GGN. To overcome the still intractable quadratic size of HCGN in P.
[Immer et al. (2021a) use structured approximations like KFAC (Martens and Grosse (013). Cheaper approximations, such as diagonal ones, often compromise accuracy (App. F Daxberger et al. 2021

---

3.4 Kronecker-Pactored Gauss-Newton Approximation (KFAC)
cronecker-factoree approximations to the Gauss-Newton, such as KFAC, are commonly used for Laplace approximations as they currently seem to provide the best known trade-off between per- formance and complexity (Ritter et al.) 2018 Daxberger al. 202T). KFAC is a block-diagonal approximation to the Gauss-Newtow matrix Haak where each block corresponds to a layer in the neural network (Martens and Grosse 2015; Botev al., 2017). It particularly efficient because each block is represented as two Kronecker factors instead of one dense matrix. For example, the GGN block of a layer with x G parameters, that is, a fully-connected layer connecting D to G neurons, would have quadratic memory complexity (D) while the corresponding Kronecker approximation KFAC is in O + G²) and is therefore even tractable for wide neural networks.
Mathematically, KFAC approximates the GGN of the ith block of the neural network parameters by enforcing a Kronecker factorisation across data points. This would only be exact for a single data point x,, for which we can write the GGN block corresponding to the parameters of the /th layer as
H = |at,n X gl. An an gl = [a,, @ [g.nAng],,) def
Ar, & G
(6)
where at, E RDD is the input to the /th layer for data point Xn. gt.n E R RGxCC is the transposed Jacobian of the network output with respect to the output of the Ith layer for Xn, and A, = (f(x,, 0)). Thus, the factors in Eq. 6 are the Jacobian terms as in the GGN (Eq. 5) but only for the Ith layer, i.e., Je, (x,)) = a & gi, KFAC then approximates the sum over N data points by summing up the Kronecker factors individually instead of breaking the Kronecker-factored structure:
H >m Iar,aan, & Ang
= 1
al.
X
Σ> m==
10.. Angi
def
where A, DixDi and G2 E RGx can be understood the uncentered covariance over N data points of the inputs to the 1th layer and the Jacobians wrt. the output of the Ith layer, respectively. The normalization by + is necessary to account for the additional terms that arise from distributing the sum over the factors. For a single-layer model, i.e., a linear model, KFAC is exact (cf. App.
4 Invariance Learning with Differentiable Laplace Approximations
We propose Laplace-GGN approximation to the marginal likelihood for invariant neural networks and enable gradient-based optimisation of their invariance parameters during training, without the use of validation data (see Fig. I for high-level overview). In our approach, we integrate the augmentation distribution p(x x. n)) into a Bayesian neural network model such that the marginal likelihood directly depends the invariance parameters 17 E H. In particular, this is due to modified likelihood function (fy 0, n), H) (c.I., Nabarro et al. 2021). This model enables optimisation of the invariance parameters using gradient ascent on the log marginal likelihood,
7 + v, log = M)),
where M H. \ {n} are remaining hyperparameters, such as regularisation strength or model architecture. However, the tractable Laplace-GGN and KFAC approximations are not available for invariant Bayesian neural networks, which prohibits straightforward application of the update in Eq. 8 using the methods described in Sec. 3
In the following, we extend the Laplace-GGr and KFAC approximations invariant neural net- works (Secs. 4. Jnd 4.2). which enables optimising the log marginal likelihood in parallel to the neural network parameters, as in [mmer al. (222Ta). However, their algorithm has an intractable memory complexity for computing the gradients w.r.t. invariance parameters or other complex hy- perparameters that act on the neural network function f directly. In practice, they only considered gradient-based optimisation of hyperparameters that act linearly on the Hessian Hg. for example regularisation strength and observation noise. In Secc. 4.3) we lift this constraint by proposing a method to obtain gradients w.r.t. complex hyperparameters without memory overhead. The final algorithm and a discussion of approximations for invariance learning are detailed in Apps. and
4.1 Laplace-GGN for Invariant Neural Network
To define the Laplace-GGN approximation to the log marginal likelihood, the hyperparameter objec- tive, we need to extend the GGN to invariant neural networks. For an invariant neural network with
5

---

parameter estimate O., the log marginal likelihood approximation is given by
logp(D) (. M) ~ Ex= logg(y,,,... (xx, 0., n),
+ log ((O. - log 2F
GOO 0_
(1)),
(9)
where the first two terms constitute the training loss of the invariant neural network corresponding to the log joint as in the vanilla Laplace approximation in Eq. 3 and the last term is the Gauss-Newton approximation High of the invariant neural network f, which we derive below. The first and last term depend on the invariance parameter 7 that we want to learn and differentiate with respect to. We use
S Monte Carlo samples EL...., Es ~ p(e) to estimate the invariant neural network as in Eq.
The log-likelihood term is approximated with S Monte Carlo samples, which leads to a lower bound.
Σ Ex=1 log logp(y, 0, M) 2 N
Be...e llggly.l EE, f(((,, 6si m); 0); M)]
N m= logp(y, E. Exi m):
(10)
which is due the concavity of the log likelihood in its natural parameter and Jensen's inequality as shown by Nabarro al.] 2021). Schwwöel et al. (2022), and detailed in App. F.T The subsequent Monte Carlo approximation is then unbiased. Increasing S leads to a tighter bound and improves the approximation. In practice, we sample 6ş independently per data point x, to reduce correlation. We can obtain stochastic gradient w.r.t. by sampling a minii-batch of M < NN data but it is not possible to batch over the S augmentations that parameterise the likelihood function. The runtime and memory complexity are therefore increased by a factor of S. For subset of data and S < 100 augmentations, the gradient w.r.t. 17 can be computed using backpropagation (Benton et al. 2020
The Gauss-Newton can be derived from the log-liillhood approximation in Eq. 10 using the same S samples. In partieular, the Jacobian and log- likelliol Hessian required for the GGN are given by
Jo(x; n) def s Jo(((x cx) 1)) and A(x; 0, m) def
c,: n):
(11)
Both terms depend on the invariance parameter v and are later differentiated with respect to it. The resulting GGN of the invariant neural network log likelihood estimated with S samples is defined as
HGGN
(n) def Σ Jx( 7 (xx,; 0, n) (x.:
(12)
The GGN for an invariant model therefore requires averaging individual Jacobians and functions of the underlying neural network f. This is in contrast to the GGN of an improper Bayesian model (Wenzel et al.) 2020) with standard data augmentation, which averages log-likelihood terms instead of functioon. In this case, entire GGN terms are averaged over the augmentations and the marginal likelihood cannot be optimised because it requires tempering (Immer et al.) 2021a). Computing the GGN of an invariant model increases the runtime by factor of over non-in con-invariant model. The empirical Fisher, which is often cheaper, can be extended analogously and requires averaging gradients instead of Jacobians. Finally, the Laplace-GGN obtained by computing the log determinant.
The Laplace-GGN approximation derived here already allows to learn small invariant neural networks with few layers and small datasets via automatic differentiation. For example, this tractable for the classification example illustrated in Fig. 1 However, computing the log determinant of the GGN has cubic complexity in the number of parameters P and is therefore intractable for deep neural networks on larger datasets. To enable its estimation, we extend KFAC to invariant neural networks in [eec. 4.2, Further, differentiating the log determinant has intractable memory complexity because does not allow for stochastic gradient but requires construction of computational graph for the entire GGN approximation. This is key limitation of the method proposed by Immer et al. (2021a) when optimising complex hyperparameters for which we provide a solution in Sec. 4.3
4.2 Extending KFAC to Invariant Neural Networks
Computing KFAC as described in Sec. 3 for invariant neural network would not preserve the Kronecker structure and therefore be intractable for deep neural networks due to the quadratic cost in the numbers of parameters per layer. In particular in En. 6, KFAC uses the fact that the Jacobian w.r.t. the parameters of the (th layer, Jo, (x,), can be written as the Kronecker product a @ gt. For an invariant neural network, the Kronecker structure cannot be maintained because of the sum over S Iugmentation-sampll Jacobians, each of which constitutes a Kronecker product:
Jo, (Xni m) E (g(x Esj E)) 3 Σ. [a & gi
(13)

---

where a , depend on the sth sample x, -((x, Eg: 77) and thus n7. In general, the sum of Kronecker products does not allow for efficient computation and requires evaluating the products. which is intractable as it requires GF) instead of O DDj + Gf) memory.
Similar the idea underlying KFAC itself (Martens and Grosse 2013), we enforce the efficient cronecker-factored structure by approximating the sum of Kronecker product as a Kronecker product of sums and appropriately normalizing by the number of terms. However, instead of applying this approximation to the GGN over multiple data points, we apply it to the Jacobian and have
Jo, (Xni و
Bl ~ 52 [
(+) & [Egh
(14)
Applying this Jacobian approximation, the with summand of the GGN can be written as
H (n)
5ラz | a
[
gi. ((x;; 0 n)
| [_.
[c)]
Σ gi.
(+)
(15)
Σ. Σ. ( @ [[E,
af (x m) E. 1 del A1, @ Gi
To compute the full KFAC, we then accumulate the Kronecker factors for the /th layer, A and Gr, of the invariant neural network over all n data points to obtain A and Gj as in Eq. 7 for vanilla KFAC. Like KFAC, our approximation for invariant models remains exact for linear models F.Ip.
The log determinant required for the nurginal-likellhovr approximation can be computed from G. A individually (Immer et al. 2021a) and has a complexity of O + GG)) as opposed to the intractable O (D This enables us to apply our method to deep invariant neural networks with widths of order 10 like vanilla KFAC. We discuss computational complexities in depth in App. D We note that KFAC for invariant neural networks could further be of independent interest for second-order optimisation and inference (Martens and Grosse 2015 Zhang et al. (0118).
4.3 Efficient Gradient Estimation of the Laplace-GGN w.r.t. Complex Hyperparameters
Automatic differentiation (AD) of the log determinant term in the Laplace-GGN w.r.t complex hyperparameters, such invariances, has an intractable memory complexity. Here, we propose a method to estimate the gradient of the log determinant in the Laplace-GGN without memory overhead. For AD, the memory complexity is equivalent to the runtime complexity of computing the log determinant, which is at least O(NP). the cost of training a deep neural network for one epoch. Such computation only tractable due batching (e.g.. for standard training losses) and otherwise would require several terabytes of memory for deep neural networks. However, the log determinant does not allow for a batched gradient and therefore AD requires storing the full training data pass.
Our approach to reducing the memory complexity relies on computing a vector-Jacobian product where both, the vector and the Jacobian, can be estimated from batches of data. Mathematically, the problem reduces to diffcrentiation of the log determinant of a sum of square matrices w.r.t. hyperparameter n, i.e., am log (T)) with H(n) = Σn. (n) E PxP. For a positive definite
matrix H, we have Blog H
H-1 Therefore, we can differentiate w.r.t. 7) with H - " (m))p-( a on 1 H.
HH
log H H(r)) Ep-
bn
vec (H- (m)) DN लy vec (H,)
(16)
where the two vectorised matrices are peediimessooal vectors and can both be computed from individual batches of H,, as follows: the first term acts as a preconditioner and can be computed by summing up the batches and inverting the resulting matrix H without storing the computation graph. The second term is a sum over N Jacobians w.r.t. 17 and we can either aggregate it or obtain an unbiased stochastic estimate from batches of data. The product of both terms constitutes vector-Jacobial product and is a standard procedure AD (Paszke et al. 2017).
The proposed method allows to aggregate gradients with respect to complex hyperparameters with memory complexity that is controlled by the batch size S M S N. contrast to naive application of AD to the log determinant, this allows to decouple memory and runtime complexity and enables gradient-based optimisation the log determinant in [q. Ww..t the augmentation parameters for deep invariant networks on large datasets. More generally, our method enables gradient-based marginal likelihood optimisation for more complex hyperparameters than previously considered (Immer et al..
202 Antorán al. 2022a). In App. E we describe the gradient computation for KFAC in detail.
7

---

learned invariance
fully-rotated MNIST
parttaly-rrottted MNIST
translated MNIST
5ca MN
regular MNIST
500
1000
500
1000
500 epochs
1000
500
1000
500
1000
epochs
epochs
epochs
epochs
horizontall transiriion
vertical translation
hotizonta scaling
verbical
sampled sugmentation:
Figure 2: Trajectories of invariance parameters v [m, 72, mo]? optimised with LILA over 1000 epochs for different versions of the MNIST dataset. From each distribution, two sampled transformations are represented as flow fields with arrow lengths scaled to pixel movement. LILA automatically learns the correct type and amount of invariance for each of the modified training datasets and keeps the other affine invariance parameters at zero as desired.
Experiments
We evaluate our method that learns invariances using Laplace approximations (LIIA) by optimising affine invariances on different MNIST (LeCun and Cortes 2010). FashionMNIST (Xiao et al. 2017). and CIFAR-10 10 (Krizhevsky al. 2009) classification tasks. To validate whether the method is capable of learning appropriate invariances, we construct several additional datasets modified by known sets of invariance transformations with the goal recover them. We consider the following affine transformations: full rotation, partial rotation, translation, and scaling (full details in App. B). We compare our approach with non-invariant baseline and Augerino (Benton et al. 2020). which is, to our knowledge, the only other method that is capable of learning invariances on complex image datasets in deep neural networks without validation data. In the non-invariant model, prior parameters are learned using the marginal likelihood (Immer et al. 2021a). For invariance learning, prior minimise parameters and 17 are jointly learned based on the marginal likelihood. For Augerino, the regular cross entropy with added regularising term -10 2|mll and a fixed weight decay of 10-4, following the original paper (Benton al.) 2020). The same parameterisation, network architecture. and initialisations (in particular = 0) were used for all methods. We assess performance of our approach by inspecting learned invariances n7. marginal likelihoods, and final test performances.
5.1 Recovering Known Invariances
To assess the invariances learned by our method LILA, we can inspect the learned invariance pa- rameters. The invariance parameter vector 7 = [m, 772 ... no]" describes affine invariances with components corresponding to x- E-translation, y-translation, rotation, horizontal and vertical scaling, and shearing (App. B). As an MLP has little symmetry encoded in the architecture itself, we expect 77 to almost correctly recover the invariances. In FFg.. 2. we plot the trajectories of each vector component over the course of training for an MLP model on different transformed MNIST datasets. As reference, we show the amount of invariance that was imposed on each dataset as dashed line. From the figure, we can observe that for each dataset, LILA learns the correct invariance as well as the amount of each invariance. To some extent, the model also learned translational invariance on the regular MNIST dataset which can be explained by intrinsic translational invariance of the dataset.
Other network architectures have certain symmetries already built-in to some extent (c.g., translational cquivariance of convolutional layers). Yet, we find that LILA also capable of inferring the correct invariances with such larger and other network architectures and on variety of datasets in App. J.3] In App. F. we further show that cheaper Hessian approximations, such as the diagonal GGN instead of KFAC and empirical Fisher (EF) instead of GGN lead to worse performance for invariance learning. This suggests that our extension of KFAC to invariant neural networks necessary for sufficient invariance learning. While Immer et al. (2021a) find that diagonal approximations can suffice for learning regularisation hyperparameters, this does not apply for the more complex invariance parameters considered here and more accurate approximations tend to increase performance.

---

fully-rotated CIFAR-10 partia lly-rotated CIFAR-10
trans ated CIFAR-10
scaled CIFAR-10
regular CIFAR-10
188 3338
accuracy
non comiimariant
30
UoAiours)
20
20
5t. 10K 50k
5k 20K 50k
1k
5k: 20k 50k
1k
5k.
10k
50k
1k
5k 10L20k 50k
subset size
Figure 3: LILA improves data efficiency on different versions of CIFAR-IO by learning invariances. We compare test accuracy on randomly sampled subsets to a non-invariant network and Augerino using& ResNet. For example on fully-rotated CIFAR-IO, invariance learning achieves the final accuracy of the non-invariant model with 10 times less data. In most cases, LILA further improves over Augerino. We plot the average performance and standard error over three seeds.
Dataset
Network
Method
Fully Rotated
Partially Rotated
Translated
Scaled
Original
-on-iinvriinnt
93.82 10 10
95,83 10.00
94,15 10.02
97.07
20.06
98.20
20.0d
MLP
Augerino
97.83 +0.03
96.35
#0.02
94.47 20.08
97.45
20.05
98.45
=0.05
MNIST
Diff. Laplace (ours)
97.74 10.07
10.11
97.28 10.00
98.33
20.05
98.98
=0.05
non-invariant
95.97 #0.33
97.31 #0.17
96.54 20.20
98,37
=0.00
99.09
=0.02
CNN
Augerino
99.04 10.02
98,91 10.03
97.79 10.00
98.77
=0.06
98.26
=0.10
LIL KFAC (ours)
98.83 +0.07
98,92 #0.05
98.69 +0.07
99.01 =0.06
99.42
=0.02
non-invariant
77.62 #0.90
81.10 #0.28
77.68
81.84 =0.05
88.48
=0.50
MLPP
Augerino
77.76
81.40
78.05 10.
82.46 20.00
89.10
20.13
F-MNIST
LILA (ours)
87.39 #0.06
86.72 #0.13
84.62 #0.08
84.31 =0.06
89.94
=0.12
non-invariant
78.69 10.25
82,12 10033
80,33 1:019
83.66
20.37
89.54
=0.20
CNN
Augerino
85.76
81.54 #0.19
82.94 +:013
83.58
=0.05
90.07
=0.12
LILA (ours)
89,45 1.0005
88.40 10.00
87.73 10.20
87.33
10.00
91.92
-0.21
non-invariant
34.16 #0.40
39.90
#0.12
69.65
16
66.06
10.15
74,13
=0.51
CIFAR-10
ResNet
Augerino
75.40 40011
74.76
#0.34
73.71
#0.31
72.07
=0.00
79.03
=1.04
LILA (ours)
79.50
77.71 10.40
79.21
10.17
76.03 10
84.19 20.76
Table 1: Test accuracy for models using LILA on different versions of the MNIST, FashionMNIST. and CIFAR- 10 datasets. We report the average performance and the standard error over three random seeds. Our method outperforms the non-invarian network and Augerino for most models and datasets.
5.2 Invariance Learning in Different Networks
Method
Test accuracy
non- invariant
85.17 +0.39
To quantify the benefit of models with learned invariances, we present final test accuracies of LILA and the baselines with different models on each of the datasets in Table I Additional marginal likelihood scores can be found in App. J.2] In terms of test accuracy and marginal like- lihood, we find that learning invariances always outper- forms the non-invariant baseline and that our approach improves over Augerino in almost all cases. This holds across the transformed and original datasets. While the improvements are modest for MNIST, the performance improvements on F-MNIST and CIFAR-10 can be up to 10 percent points. In Table the commonly used Wide ResNet architecture on CIFAR-10 and find that invariance learning with Augerino merely improves performance while our method achieves performance improvement of almost 7% points. We also report the performance of LILA with the cheaper KFAC-EF instead of GGN.
Augerino
87.67 #0.08
LILA (ours)
91.98 10.04
LILA EF (ours)
91.64 10.26
Table 2: Invariance Learning with KFAC on CIFAR-10 with a Wide ResNet achieves best test accuracy. Also, the cheaper empirical Fisher (EF) variant im- proves over Augerino.
we use
5.3 Invariance Learning Improves Data Efficiency
Invariances can be particularly useful for data-efficiency, which can be evaluated by measuring performance on subsets of data. In Fig. 3] we show the test accuracy of LILA, the non-invariant baseline, and Augerino trained on different subsets of CIFAR-IO and its modified versions. We further provide results for all architectures and datasets in App. J.5) In general, we find that invariance learning with both Augerino and our method always improves performance across datasets and subset sizes. Most notably on a subset of 1000 regular 10 dataset samples, Augerino improved performance by 18 percentage points compared to the non-invariant baseline, whereas our method showed an improvement of 22 percentage points. LIIA requires only 10% of the data to obtain the same accuracy as the non-invariant model on the fulle-rotated dataset. These findings suggest that learning invariances is useful in general, and in particular when limited data is available.

---

5.4 On the Learned Distributions
In some instances (see CIFAR-10 results Apps. J.3 and J.4p. we observed that the model learned translational invariance on the scaled dataset rather than scale invariance. Our independent uniform distributions over invariance components cannot capture correlations between components whereas the scaled dataset was jointly scaled across the horizontal and vertical axes and thus is correlated. We hypothesise that more complex distributions that do allow for capturing correlations could offer a potential solution in such cases, but leave investigation of more complex families to future work.
6 Discussion and Limitations
We discuss the runtime complexity and approximations our method in general, and for invariance learning in particular, in detail in [App. D and [App. F} respectively. The runtime complexity of LILA, like that of Augerino and other methods that use test-time data augmentation, increases linearly by the factor S, which denotes the number of augmentation samples used. Since we know that we are sampling from lower bound (Eq. 24), more samples are generally better and are expected to improve the performance, which we also observed in our experiments. In addition, LILA requires estimation and differentiation of the log-determinant term. While our extension of KFAC to sampling-base based invariant models and the batched gradients make such computation at all tractable, it can still be expensive for a large number of classes C' and augmentation samples Using the EF instead of the GGN overcomes scaling in C' but is cruder approximation. Although these additional computations make LILA slower than Augerino, they enable overcoming its issues, i... parameterisationndepppe deppnddeeeeeeeeeeeeeeeeeeeeeeeee and additional hyperparameters that need to be tuned (App. C, enable to learn soft invariances much more preciselly.
Conclusion
We presented method that enables automatic invariance learning in deep neural networks directly from training data, without requiring supervision or validation data. The approach inspired by using the marginal likelihood, which is irrameteriatioo-indeeeedddd generalisation performance.
quantity coinciding with make this practical, we use differentiable Laplace approximation to allow for gradient-based optimisation of invariances in deep learning. While the accuracy of the approximation is difficult to verify, we do show experimentally that the method is capable of learning invariances in MNIST, FaSHionMNIST, and CIFAR- 10 datasets, leading to better marginal likelihoods and higher test performances. Our work shows that approximate Bayesian inference methods can be useful for learning complex hyperparameters, even in deep learning, and are therefore relevant beyond predictive uncertainty estimation. In future work, it would be interesting to improve the scalability and accuracy of marginal likelihood approximations which could enable learning even more complex hyperparameters, such as augmentation distributions parameterised by neural networks. Alternatively, improving parameterisations of invariances in neural networks could greatly improve the scalability of LILA and related approaches.
Acknowledgements
A.I. acknowledges funding by the Max Planck ETH Center for Learning Systems (CLS). V.F. acknowledges funding by the Swiss Data Science Center through a PhD Fellowship. the Swiss National Science Foundation through Postdoc.Mobilii Fellowship, St John's 's College Cambridge through a Research Fellowship. and the Branco Weiss Foundation through a Branco Weiss Fellowship.
References
Javier Antorán, Riccardo Barbano, Johannes Leuschner, José Miguel Hernández-Lobato and Bangti Jin. A probabilistic deep image prior for computational tomography. arXiv preprint 000.002003.00479. 2022a.
Javier Antorán, David Janz, James U Allingham, Erik Daxberger, Riccardo Rb Barbano, Eric Nalisnick, and José Miguel Hernández-Lobato. Adapting the linearised laplace model evidence for modern deep learning, In International Conference on Machine Learning, 2022b.
10

---

Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E Smidt, and Boris Kozinsky. Se equiviarai graph neural networks for data-efficient and accurate interatomic potentials. arXiv preprint arXiv:2101.00 2044003644 2021.
Gregory Benton, Marc Finzi, Pavel Izmailov, and Andrew Gordon Wilson. Learning invariances in neural networks. arXiv preprint OrXiv:2010.11880 2020.
Christopher M Bishop. Pattern recognition and machine learning. Information Science and Statistics. Springer, 2006.
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural networks. In Proceedings of the 32nd International Conference on Machine Learning, pages 1613-1622, 2015.
Aleksandar Botev, Hippolyt Ritter, and David Barber. Practical Gauss-Newton optimisation for deep learning. In International Conference Machine Learning, International Convention Centre, Sydney, Australia, 2017. PMLR.
Léon Bottou. Lcage-scale machine learning with stochastic gradient descent. In Proceedings of COMPSTAT pages 177-186. Springer, 2010.
Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik Bekkers, and Max Welling. Ge- ometric and physical quantities improve (3) equivariant message passing. arXiv preprint arXiv:2110.02905. 2021.
Taco Cohen and Max Welling. Group equivariant convolutional networks. In International conference on machine learning, pages 2990-2999. PMLR, 2016.
Taco S Cohen, Mario Geiger, Jonas Köhler, and Max Welling, Spherical cnns. arXiv preprint 10X1v01801.10110, 2018
Ekin D Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoe V Le. Autoaugment: 2018.
Learning augmentation policies from data. arXiv preprint arXiv:
Felix Dangel, Frederik Kunstner, and Philipp Hennig. Backpack: Packing more into backprop. In Proceedings of 7th International Conference on Learning Representations, 2019.
Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig. Laplace redux-effortless bayesian deep learning. Advances in Neural Information Processing Systems, 34, 2021.
Guillaume P Dehaene. A deterministic and computable bernstein-von mises theorem, arXiv preprint 904.05990..02505, 2019.
Edwin Fong and CC Holmes. On the marginal likelihood and cross-validation. Biometrika, 107(2): 489-496. 2020.
Andrew YK Foong, Yingzhen Li, José Miguel Hernández-Looato, and Richard E Turner. 'in- between' uncertainty in bayesian neural networks. arXiv preprint arXiv:/906 /537, 2019.
Kunihiko Fukushima and Sei Miyake. Neocognitron: A self-organizing neural network model for mechanism of visual pattern recognition. In Competition and cooperation neural nets, pages 267-285. Springer, 1982.
Pascal Germain, Francis Bach, Alexandre Lacoste, and Simon Lacoste-Julien. Pac-bayesian theory meets bayesian inference. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 29, Curran Associates, Inc., 2016.
David Ginsbourger, Olivier Roustant, and Nicolas Durrande. On degeneracy and invariances of random fields paths with applications in gaussian process modelling. Journal of Statistical Planning and Inference, 170:117-128, 2016. ISSN 0378-3758.
Peter D Grünwald. The minimum description length principle. MIT press, 2007.

---

Kaiming He. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition. pages 770-778, 2016.
Alexander Immer, Matthias Bauer, Vincent Fortuin, Gunnar Rätsch, and Mohammad Emtiyaz Khan. Scalable marginal likelihood estimation for model selection in deep learning. arXiv preprint 1xXiv:2104.04975. 2021a.
Alexander Immer, Maciej Korzepa, and Matthias Bauer. Improving predictions of bayesian neural nets via local lincarization. In Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, pages 703-711, 2021b.
Alexander Immer, Lucas Torroba Hennigen, Vincent Fortuin, and Ryan Cotterell. Probing as quantifying inductive bias. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics, pages 1839-1851,200 2022. 2022.
1851,
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in neural information processing systems, pages 8571-8580, 2018.
Max Jaderberg. Karen Simonyan, Andrew Zisserman, et al. Spatial transformer networks. Advances in neural information processing systems, 28:2017-20255, 2015.
Mohammad Emtiyaz. E Khan, Alexander Immer, Ehsan Abedi, and Maciej Korzepa. Approximate inference turns deep networks into gaussian processes. In Advances in Neural Information Processing Systems, pages 3088-3098, 2019.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015.
Diederik Kingma and Max Welling. Auto-encoding variational bayes.
arXiv preprint
arXiv:1312.6114, 2013.
Imre Risi Kondor. Group theoretical methods in machine learning. 2008
Risi Kondor, Zhen Lin, and Shubhendu Trivedi. flebsch-gordan nets: fully fourier space spherical convolutional neural network. Advances in Neural Information Processing Systems, 31:10117- 10126, 2018.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Frederik Kunstner, Philipp Hennig, and Lukas Balles. Limitations of the empirical fisher approxima- tion for natural gradient descent. In Advances Neural Information Processing Systems, pages 4168-4169, 2019.
Pierre-Simon de Laplace. Mémoire sur la probabilité des causes par les événements. Mémoires de l'Académie royale des sciences de Paris (Savants étrangers), 6:621-656, 1774.
Yann LeCun and Corinna Cortes. http://aann.eeu..coo/exxxbdii
MNIST handwritten digit database.
2010. URL http:://tannn.....eeee/ lecun. www.cuncooo/eee httom/exdm/mnieee
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11))2278-2322 1998.
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 Intelligncee and Statistics, volume 108 of Proceedings of Machine Learning Research, pages 1540-1552. PMLR, 26-28 Aug 2020.
David JC MacKay. A practical bayesian framework for backpropagation networks. Neural computa- tion, 4(3):4484477 1992.
David JC MacKay. Information theory, inference and learning algorithms. Cambridge university press, 2003.
12

---

James Martens. New insights and perspectives on the natural gradient method. Journal of Machine Learning Research, 21 21(146):1-7,, 2020.
James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curvature. In International conference on machine learning, pages 2408-2417, 2015.
Cleve Moler and Charles Van Loan. Nineteen dubious ways to compute the exponential of a matrix, twenty-five years later. SIAM review, 45(1):3 2003.
Kevin P Murphy. Machine learning: a probabilistic perspective. MIT press, 2012.
Seth Nabarro, Stoil Ganev, Adrià Garriga-Alonso, Vincent Fortuin, Mark van der Wilk, and Laurence Aitchison. Data augmentation bayesian neural networks and the cold posterior effect. arXiv preprint wrXiv:2106.055%6 2021.
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
PMLR, 2021.
Sebastian Ober, Carl E. Rasmussen, and Mark van der Wilk. The promises and pitfalls of deep kernel learning. In Cassio de Campos and Marloes H. Maathuis, editors, Proceedings of the Thirty- Seventh Conference on Uncertainty in Artificial Intelligence (UAI), volume 161 of Proceedings of Machine Learning Research, pages 1206-1216. PMLR, 27-30 Jul 2021.
Kazuki Osawa, Siddharth Swaroop, Mohammad Emtiyaz Khan, Anirudh Jain, Runa Eschenhagen, Richard E Turner, and Rio Yokota. Practical deep learning with bayesian principles. In Advances in Neural Information Processing Systems, pages 4289 4289-4301, 2019.
Kazuki Osawa, Yohei Tsuji, Yuichiro Ueno, Akira Naruse, Chuan-Shena Foo, and Rio Yokota. Scalable and practical natural gradient for large-scale deep learning, IEEE Transactions on Pattern Analysis and Machine Intelligence, 2020.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang. Zachary Vetooo, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
Carl Edward Rasmussen and Zoubin Ghahramani. Occam's razor. In Advances neural information processing systems, pages 294-300, 2001.
Carl Edward Rasmussen and Christopher KI Williams. Gaussian processes for machine learning. MIT press Cambridge, MA, 2006.
Hippolyt Ritter, Aleksandar Botev, and David Barber. scalable laplace approximation for neural networks. In International Conference on Learning Representations, 2018.
Pola Schwöbel, Martin Jørgensen, Sebastian W. Ober, and Mark van der Wilk. Last layer marginal likelihood for invariance learning. In Proceedings of the Twenty Fifth International Conference on Artificial Intelligence and Statistics (AISTATS), 2022.
Joost van Amersfoort, Lewis Smith, Andrew Jesson, Oscar Key, and Yarin Gal. Improving determinis- tic uncertainty estimation in deep learning for classification and regression. CoRR, abs/2102.1 11409, 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.
Elise van der Pol, Daniel Worrall, Herke van Hoof, Frans Oliehock, and Max Welling. Mdp homomorphic networks: Group symmetries in reinforcement learning. Advances in Neural Information Processing Systems, 33, 2020.
Mark van der Wilk, Matthias Bauer, ST John, and James Hensman. Learning invariances using the marginal likelihood. Advances Neural Information Processing Systems, pages 1938-9948, 2018.
13

---

Florian Wenzel, Kevin Roth, Bastiaan S Veeling, Jakub Swiatkowski, Linh Tran, Stephan Mandt, Jasper Snock, Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin. How good the bayes posterior in deep neural networks really? In International Conference Machine Learning, 2020.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms, 2017.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In BMVC. BMVA Press, 2016.
Guodong Zhang, Shengyang Sun, David Duvenaud, and Roger Grosse. Noisy natural gradient as variational inference. In International Conference on Machine Learning, pages 5852-5861, 2018.
Hongyi Zhang, Yann N Dauphin, and Tengyu Ma. Fixup initialization: Residual learning without normalization. In International Conference on Learning Representations, 2019.
Allan Zhou, Tom Knowles, and Chelsea Finn. Meta-learning symmetries by reparameterization. arXiv preprint urXiv:2007.02933, 2020.
14

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations your work? [Yes] We point out all approximations
and discuss complexities of the method.
(c) Did you discuss any potential negative societal impacts of your work? [N/A] There are no direct societal impacts.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full of assumptions all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)? [Yes] Code available com/tychovdo/111 /ila Details are in Sec. and App.
on hattps://githue
(b) Did you specify all the training details (c.g.. data splits, hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g.. with respect the random seed after running experi- ments multiple times)? [Yes] all experiments are repeated times to estimate the error bars.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [No]
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [No]
(c) Did you include any new assets either the supplemental material or as URL? [No]
(d) Did you discuss whether and how consent was obtained from people whose data you'l 're using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
15


Reviews:
[{"review_id": "j4qpPSoeuob", "reviewer": "Reviewer_YcvR", "summary": "The paper proposes a gradient-based method for automatically selecting the data augmentation parameters without the validation data. The trick is to parameterize the data augmentation scheme, approximate the marginal likelihood with respect to these data augmentation parameters and optimize them during training. The authors propose schemes such as KFAC approximation and explicit automatic differentiation to make the computation of differentiable marginal likelihood more tractable. Empirically, the proposed approach is capable of recovering invariances presented in the data on MNIST, FashionMNIST, and CIFAR10.", "questions": "Comments:\n* In the introduction, the authors state that choosing the correct transformation is expensive and requires trial and error. However, the proposed approach also requires selecting the parameterization of $\\eta$, which also requires some prior knowledge of the dataset and task. For example, in the case of image classification, affine invariance parameterization looks like a reasonable choice, but how should we select the parameterization for other tasks (e.g., language modelling). It would be helpful if the authors addressed these issues.\n* In Table 2, is the model evaluated on the standard CIFAR-10 datasets, or is it a corrupted version (e.g., scaled and translated)? The test accuracy for the baseline method seems to be a little lower than I expected.\n* As the paper introduces a lot of notations, it would be helpful to have a table of notations summarizing all notations used in the paper in the Appendix.\n* As the authors mentioned in the checklist, releasing the code would be helpful.\n* What additional hyperparameters does the proposed method require, and how sensitive are they? Given a new dataset, can the proposed algorithm work without putting a lot of effort into tuning these hyperparameters?\n\nMinor Comments:\n* Some colours in Figure 2 (e.g. vertical & horizontal scaling) are too similar, and it is difficult to distinguish them in the plot when printed.\n* In line 301, there is a missing bracket after App. B.", "limitations": "The work does not have a potential negative societal impact. ", "rating": 7, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* The topic is interesting and relevant to the Neurips community.\n* The paper is well-written, and the technical details look correct to me. \n* While limited to small datasets and models, the experiments demonstrate that the proposed approach can successfully recover invariances presented in the model and can help improve generalization.\n\nWeaknesses:\n* There are several specific limitations of the proposed approach. To give some specific examples: (1) the proposed approach does not scale to large datasets and models, (2) the method requires task-specific data augmentation parameterization (e.g., affine transformation for image classification tasks), and it is not straightforward to apply the proposed method in other domains, and (3) the proposed algorithm is challenging to implement on standard deep learning frameworks and introduce additional computational overhead that scales quadratically with model parameters and linearly to the number of datasets.\n* While I believe that the toy experiments presented in the paper demonstrate the proposed approach's effectiveness, the proposed method's scalability is still questionable.\n* It would be helpful if the authors also compared their algorithms with other data augmentation adaptation methods that make use of validation datasets. To my understanding, these methods should be much more efficient than the proposed approach, although it requires a validation dataset. If we select 20% of training data as validation data, shouldn't the prior algorithms also be able to learn these invariances?\n\n\n** Post-Rebuttal **\n\nThank you for the author's response. I acknowledge that I have read the rebuttal and other reviewers' comments. As the author addressed all my concerns in the paper, I increased my score.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "TD7pbBhQoF", "reviewer": "Reviewer_LhTL", "summary": "The authors extend previously proposed techniques for a Laplace approximation for Bayesian model selection to learn and incorporate invariances while training a deep network. To make this approximation computationally feasible, the authors propose techniques to apply a Kronecker-Factored approximation. The authors extensively evaluate their method on many datasets, with ablations and comparisons to another method that learns invariances while training. ", "questions": "* Does the baseline get trained with any augmentations applied?\n* Have the authors considered using a generative model in place of the six hand-picked invariances (e.g., a disentangled variational auto-encoder like the Beta-VAE from Higgins et al 2011)? This would take the work one step closer to true a priori invariance learning. ", "limitations": "* From the experiments the authors have presented, it appears that their methods requires the set of candidate invariances to be specified beforehand. This is limited for datasets where the invariances are not known beforehand.\n* The authors have not released code.", "rating": 7, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n* The authors have very illustrative experiments (Figure 2 is very well-presented) that show that the models they train are learning the invariances present in the dataset. Training a network to capture the invariances present in a dataset is of significant interest for the machine learning community. Furthermore, they test their method on a wide variety of datasets/network architectures. \n* The authors clearly describe the novelty of their approach in comparison to Immer et al 2021. They describe how the aforementioned work would be computationally infeasible for their proposed task.\n\nWeaknesses\n* One goal of this work is “obtain a model that is invariant to a set or distribution of transformations” (lines 16-17) as measured by test accuracy in Table 1. In the case where one wants s network with invariance to the natural variations in a dataset, the authors should justify their approach over simply training with augmented samples, as done in (RandAugment, Cubuk et al. 2019). For example, can the authors’ approach provide benefit over RandAugment on ImageNet using their six candidate augmentations? It should be stated that the authors’ proposed method has significance outside this scenario, since they also identify which invariances are present (as opposed to augmentation approaches). Nevertheless, such a comparison is important for the community to understand which method should be used for this given circumstance.\n* Figure 4 of the appendix is very nice for measuring performance of different approximation methods. To strengthen their work, the authors should consider giving empirical runtime results as well so that practitioners can understand the empirical performance/run speed tradeoffs. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "8YsA-g2gycJ", "reviewer": "Reviewer_Dhxp", "summary": "The paper proposes to select parameters for augmentation through a training procedure (without validation data, like the prior work on Augerino). It does so by introducing a Bayesian model selection framework for selecting the ideal set of augmentations for a neural network. Effort in making this model tractable is notable and the evaluation shows some benefits over strong baselines.", "questions": "* What parts of the design of your model are entirely novel and what parts are adapted from prior work (e.g., how KFAC, Laplace-GGN, etc. fit in with prior work)?\n* What are the computational concerns with this method? How does it compare with Augerino in terms of training and inference time?\n* What is the significance of the baseline used in experiments?", "limitations": "N/A", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* The paper is clear and reads well, though the treatment of technical derivations is quite long.\n* Model does sensible behavior (i.e., Figure 2).\n* Model has improvements over Augerino (e.g., Table 1).\n* Diversity of datasets, models, and augmentations seems acceptable, but can be still strengthened (e.g., Table 1).\n\nWeaknesses:\n* One baseline is non-invariant, which seems like a weak baseline, since there is often a significant performance boost by just using augmentations in a simple way. Augerino is a stronger baseline, but the results are much less noticable against that baseline (e.g., Figure 3).\n* It's not clear nor intuitive why this method works. Some examples of failure modes of Augerino that were addressed would be enlightening.\n\nPost Rebuttal:\nMany of the paper's strengths are due to being a Bayesian treatment of the problem of setting augmentation parameters, which avoids hyperparameter tuning. There is a convincing case that it would be used in the future as a building block as this approach works on a standard set of benchmarks and is otherwise sensible. The evaluation can be strengthened (as always, e.g., scaling the approach on more difficult datasets), but I think the paper has sufficient contribution as is, and so I am increasing my score.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept