SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Tycho F.A. van der Ouderaa<sup>1</sup>
Alexander Immer2,3
Mark van der Wilk<sup>1,4</sup>
$^{1}$ Department of Computing, Imperial College London, United Kingdom
$^{2}$ Department of Computer Science, ETH Zurich, Switzerland
3Max Planck Institute for Intelligent Systems, Tübingen, Germany
$^{4}$ Department of Computer Science, University of Oxford, United Kingdom
Convolutions encode equivariance symmetries into neural networks leading to better generalisation performance. However, symmetries provide fixed hard constraints on the functions a network can represent, need to be specified in advance, and can not be adapted. Our goal is to allow flexible symmetry constraints that can automatically be learned from data using gradients. Learning symmetry and associated weight connectivity structures from scratch is difficult for two reasons. First, it requires efficient and flexible parameterisations of layer-wise equivariances. Secondly, symmetries act as constraints and are therefore not encouraged by training losses measuring data fit. To overcome these challenges, we improve parameterisations of soft equivariance and learn the amount of equivariance in layers by optimising the marginal likelihood, estimated using differentiable Laplace approximations. The objective balances data fit and model complexity enabling layer-wise symmetry discovery in deep networks. We demonstrate the ability to automatically learn layer-wise equivariances on image classification tasks, achieving equivalent or improved performance over baselines with hard-coded symmetry.
Symmetry constraints, such as layer-wise equivariances of convolutional layers, allow neural networks to generalise efficiently [Cohen and Welling, 2016]. However, it is not always clear which symmetries to use and how much they should be enforced. In digit recognition, some rotational robustness is desirable, but strict invariance may make 6's and 9's indistinguishable. Similarly, translation equivariance prevents rate coding of relevant positional information [Sabour et al., 2017]. We propose a method to automatically learn the type and amount of layer-wise equivariances from training data.
To do so, we relax symmetries building upon recent literature on approximate equivariance, including residual pathways [Finzi et al., 2021a] and non-stationary filters [van der Ouderaa et al., 2022], effectively allowing controllable interpolation between strict equivariance or less constrained functions. To remain practical, we propose extensions to keep the parameter count similar to classic convolutions.
Allowing for differentiable selection between equivariances alone is not enough, as symmetries act as constraints and are therefore not encouraged by training losses that measure data fit, as noted in van der Ouderaa and van der Wilk [2021], Immer et al. [2021]. Search over many continuous hyperparameters with cross-validation is infeasible, relying on expensive retraining, hold-out validation data and does not leverage gradients. Instead, we follow Bayesian model selection through differentiable Laplace approximations to learn layer-wise equivariances from training data in a single training procedure.
We demonstrate automatically learning layer-wise symmetry structure on image classification tasks. To do so, we improve upon existing parameterisations of differentiable equivariance and derive corresponding Kronecker-factored Laplace approximations to the marginal likelihood. On image classification, we show that our method automatically learns convolutional structure in early layers and achieves similar or improved performance compared to architectures with hard-coded symmetry.
Code accompanying this paper is available at https://github.com/tychovdo/ella37th Conference on Neural Information Processing Systems (NeurIPS 2023).
Enforcing strict equivariance into the architecture. It can be shown that enforcing equivariance constraints in linear neural network layers is equivalent to performing a convolution [Kondor and Trivedi, 2018, Cohen et al., 2019]. Regular convolutional layers embed equivariance to translation, but this notion can be extended to other groups by convolving over groups [Cohen and Welling, 2016]. Group convolutions have been proposed for various domains and more complex continuous groups [Weiler et al., 2018, Weiler and Cesa, 2019, Worrall et al., 2017]. Equivariances enforce hard constraints on functions and can even form a linear subspace that can be explicitly computed [Van der Pol et al., 2020, Finzi et al., 2021b].
Place coding, rate coding, disentanglement and part-whole hierarchies. Equivariance closely relates to disentangled representations between place-coded and rate-coded features found in capsule networks [Sabour et al., 2017, Hinton et al., 2018, Kosiorek et al., 2019], as formalised in Cohen and Welling [2014], Cohen et al. [2018]. Positional information under strict translation equivariance of convolutions can only be place-coded, not rate-coded in features. Capsule networks aim to break this strict disentanglement to allow the embedding of part-whole hierarchies [Sabour et al., 2017]. We argue that our method solves the same problem by relaxing the disentanglement of strict equivariance and provide an automatic procedure to learn the extent to which information is exchanged and moved from place-coded to rate-coded features.
Approximate equivariance. Symmetries are often assumed known, must be fixed, and can not be adapted. This can be desirable if symmetries are known [Veeling et al., 2018], but becomes overly restrictive in other real-world tasks where symmetries or the appropriate amount of symmetry is not clear. Enforcing symmetry regardless can hamper performance [Liu et al., 2018b, Wang et al., 2022].
Approximate symmetry can prevent symmetry misspecification without losing the benefits that inductive biases of symmetries provide. Although strict group symmetries always apply to the entire group, by the group theoretical axiom of closure, approximate symmetry can be thought of as a form of robustness towards group actions that is only locally enforced (around training data). We use relaxations of symmetry that follow notions of approximate symmetry in literature [van der Wilk et al., 2018, Wang et al., 2020]. Although invariances can often easily be relaxed by simply averaging data augmentations [Benton et al., 2020, Schwöbel et al., 2021, van der Ouderaa and van der Wilk, 2021], parameterising approximate equivariance [Wang et al., 2022] can be more difficult. We build upon recent attempts, including partial equivariance [Romero and Lohit, 2021], residual pathways [Finzi et al., 2021a], and non-stationary filters [van der Ouderaa et al., 2022, van der Ouderaa and van der Wilk, 2023]. Existing relaxations of equivariance often introduce many additional parameters due to added linear layers [Finzi et al., 2021a] or use of hyper networks [Romero and Lohit, 2021, van der Ouderaa et al., 2022], making adoption in practice difficult. We strive to keep parameter counts close to classical convolutional layers, sometimes requiring further simplifying assumptions.
Symmetry discovery and objectives To learn symmetry, we relax layer-wise symmetries and thereby allow them to be differentiably learned with gradients, with strict symmetry as limiting case. Since adding symmetries does not directly improve training losses that rely on training fit, there is no direct encouragement of the objective function to use symmetries. To overcome this, some works have considered using validation data [Maile et al., 2022, Zhou et al., 2020]. Alternatively, Benton et al. [2020], Finzi et al. [2021a], Yang et al. [2023], van der Ouderaa et al. [2022] consider learning symmetry from training data only, but to do so require explicit regularisation that needs tuning. Some problems with this approach have been noted in Immer et al. [2022], which, like van der Wilk et al. [2018], proposes to use Bayesian model selection to learn invariances in a single training procedure, use differentiable Laplace approximations to remain tractable in deep neural networks. The method only considers invariances, which can be less efficient because it relies on learning parameters in the sampler instead of learning regularized symmetry in intermediary representations. We extend this line of work to enable automatic learning of layer-wise equivariances from training data.
Neural architecture search As layer-wise equivariances imply a convolutional architectural structure [Cohen et al., 2019], symmetry discovery can be viewed as a form of neural architecture search (NAS). Our approach uses relaxed symmetries to allow differentiable search over architectural structures, similar to DARTS [Liu et al., 2018a]. Instead of using validation data, however, our method uses Bayesian model selection to learn architecture from training data only.
We distinguish ourselves from large parts of AutoML literature in that we do not require expensive outer loops [Zoph et al., 2018], external agents [Cubuk et al., 2018], or even validation data [Lorraine et al., 2020, Yeh et al., 2022]. Our method can be interpreted as a one-shot architecture search using Laplace approximations, similar to [Zhou et al., 2019]. We demonstrate equivariance learning, but see no theoretical reason why the proposed objective should be limited to symmetry discovery, making it potentially broadly applicable in architecture search. For instance, in approaches that aim to learn network depth [Antorán et al., 2020], filter sizes [Romero and Zeghidour, 2023] and other layer types from training data. Although interesting future work, we focus on learning layer-wise equivariances.
Bayesian Model Selection We propose to treat layer-wise equivariances as hyperparameters and learn them with Bayesian model selection – a well-understood statistical procedure. Traditionally, model selection has only been available for models with tractable marginal likelihoods, such as Gaussian Processes [Rasmussen, 2003], but approximations for deep learning have also been developed, such as Variational Inference schemes [Ober and Aitchison, 2021], deep Gaussian Processes [Dutordoir et al., 2020], and scalable Laplace approximations [Immer et al., 2021], which we use. Although application in deep learning involves additional approximations, the underlying mechanism is similar to classic Automatic Relevance Determination [MacKay et al., 1994] to optimise exponential basis function lengthscales. Even more so, as we cast layer-wise equivariance learning as a lengthscale selection problem in Section 3.2. Learning invariances with Bayesian model selection in GPs was demonstrated in van der Wilk et al. [2018], and extensions to neural networks exist [Schwöbel et al., 2021, van der Ouderaa and van der Wilk, 2021], but often use lower bounds that only work on shallow networks. For deep networks, Immer et al. [2022, 2023] and Mlodozeniec et al. [2023] recently demonstrated invariance learning with marginal likelihood estimates in large ResNets.
This work describes a method that enables layer-wise learning of equivariance symmetries from training data. We formulate different priors consisting of hyperparameters that place a different mass on functions that obey equivariance symmetry constraints and functions that do not. The amount of symmetry becomes a hyperparameter we empirically learn from data using Bayesian model selection.
In Section 3, we discuss how to efficiently parameterise relaxed equivariance in convolutional architectures. We propose additional factorisations and sparsifications in Section 3.1 to bring the parameter count of relaxed equivariance closer to that of classical convolutions. In Section 3.2, we describe how the parameterisations can be used to specify the amount of equivariance in the prior.
In Section 4, we deal with the objective function and discuss approximate marginal likelihood for layer-wise relaxed equivariance. We rely on recent advancements in scalable linearised Laplace approximations that use KFAC approximations of the Hessian, and in Sections 5.1 and 5.2 derive KFAC for proposed relaxed equivariance layers to use them in conjunction with this objective.
3 Differentiable parameterisation of layer-wise equivariance symmetry
To learn symmetries from training data, we use the strategy of relaxing equivariance constraints through parameterisations allowing interpolation between general linear mappings and linear mappings that are constrained to be strictly equivariant. In our choice of relaxation, we strive to keep parameter counts as close as possible to classical convolutional layers to remain useful in practice. We describe parameterisations of linear, equivariant and relaxed equivariant layers for discrete 2-dimensional translation symmetry  $\mathbb{Z}^2$ , as this is the most widely used symmetry in deep learning. All layers in this work can be naturally extended to other symmetry groups, as described in App. A.
Fully-connected layer A fully-connected layer forms a linear map between an input feature map  $\mathbf{x}:[0,C]\times \mathbb{Z}^2\to \mathbb{R}$  (spatially supported at intervals  $[0,X]\times [0,Y])$  , with  $C$  input channels outputting a  $C^\prime$  channel feature map  $\mathbf{y}:[0,C^{\prime}]\times \mathbb{Z}^{2}\rightarrow \mathbb{R}$  with the same spatial support:
$$
\mathbf {y} \left(c ^ {\prime}, x ^ {\prime}, y ^ {\prime}\right) = \sum_ {c} \sum_ {x, y} \mathbf {x} (c, x, y) \mathbf {w} \left(c ^ {\prime}, c, x ^ {\prime}, y ^ {\prime}, x, y\right) \tag {1}
$$
where non-zero weights in  $\mathbf{w}:\mathbb{Z}^6\mapsto \mathbb{R}$  can be represented in  $C^\prime CX^2 Y^2$  parameters. This may be more clear if, instead of using real-valued functions to describe features and weights, we equivalently flatten  $\mathbf{x},\mathbf{y}$  as vectors  $\pmb {y} = \pmb {W}\pmb{x}$  and write  $\mathbf{w}$  as a matrix  $\pmb {W}\in \mathbb{R}^{C'XY\times CXY}$ . Yet, we stick to the form of Eq. (1) as this simplifies future factorisations and generalisations. Furthermore, the notation emphasises that other layers, such as convolutions, form special cases of a fully-connected layer.
Convolutional layer We might want to constrain a layer to be strict equivariant, so that translations in the input  $(dx,dy)\in \mathbb{Z}^2$  yield equivalent translations in the output, mathematically:
$$
\mathbf {y} \left(c ^ {\prime}, x ^ {\prime} + d x, y ^ {\prime} + d y\right) = \sum_ {c} \sum_ {x, y} \mathbf {x} \left(c, x + d x, y + d y\right) \mathbf {w} \left(c ^ {\prime}, c, x ^ {\prime}, y ^ {\prime}, x, y\right) \quad \forall d x, d y \in \mathbb {Z} \tag {2}
$$
It can be shown that equivariance is equivalent (if and only if) to having a weight matrix that is stationary along the spatial dimension, meaning it solely depends on relative coordinates  $\bar{x} = x' - x$  and  $\bar{y} = y' - y$  and can thus be written in terms of stationary weights  $\bar{\theta}:(c,c',\bar{x},\bar{y})\mapsto \mathbb{R}$ :
$$
\mathbf {y} \left(c ^ {\prime}, x ^ {\prime}, y ^ {\prime}\right) = \sum_ {c} \sum_ {x, y} \mathbf {x} (c, x, y) \bar {\boldsymbol {\theta}} \left(c ^ {\prime}, c, x ^ {\prime} - x, y ^ {\prime} - y\right), \tag {3}
$$
with  $C^\prime CXY$  parameters in  $\bar{\theta}$ . Note that Eq. (3) is exactly a convolutional layer ('convolution is all you need', Theorem 3.1 of Cohen et al. [2019]). Constraining the stationary weights  $\bar{\theta}$  to be only locally supported on a rectangular domain  $\bar{x},\bar{y}\in [-\frac{S - 1}{2},\frac{S - 1}{2}]$ , and 0 otherwise, gives rise to small (e.g.  $3\times 3$ ) convolutional filters and further reduces the required parameters to  $C^\prime CS^2$ .
3.1 Factoring and sparsifying relaxed equivariance
To learn layer-wise equivariances, we take the approach of relaxing symmetry constraints in neural network layers allowing explicit differentiable interpolation between equivariant and non-equivariant solutions. As a starting point, we consider the residual pathways of [Finzi et al., 2021a], reinterpreted as factorisation of the linear map as  $\mathbf{w} = \pmb{\theta} + \bar{\pmb{\theta}}$  with fully-connected  $\pmb{\theta}: (c', c, x', y', x, y) \mapsto \mathbb{R}$  and stationary  $\bar{\pmb{\theta}}(c', c, \bar{x}, \bar{y}) \mapsto \mathbb{R}$  with  $\bar{x} = x' - x$  and  $\bar{y} = y' - y$ :
$$
\mathbf {y} \left(c ^ {\prime}, x ^ {\prime}, y ^ {\prime}\right) = \underbrace {\sum_ {c} \sum_ {x , y} \mathbf {x} \left(c , x , y\right) \boldsymbol {\theta} \left(c ^ {\prime} , c , x ^ {\prime} , y ^ {\prime} , x , y\right)} _ {\text {f u l l y - c o n n e c t e d (F C)}} + \underbrace {\sum_ {c} \sum_ {x , y} \mathbf {x} \left(c , x , y\right) \bar {\boldsymbol {\theta}} \left(c ^ {\prime} , c , x ^ {\prime} - x , y ^ {\prime} - y\right)} _ {\text {c o n v o l u t i o n a l l a y e r (C O N V)}} \tag {4}
$$
The layer consists of simply adding FC and CONV layer outputs. The CONV path spans the linear subspace of FC that obeys strict equivariance of Eq. (2). As such, the combined layer is equally expressive as a linear map (we can construct a bijection between  $\mathbf{w}$  and  $\theta +\bar{\theta}$ , as is common for relaxed equivariance [van der Ouderaa et al., 2022]). In Finzi et al. [2021a], independent priors are placed on both sets of weights with relative variances controlling how much equivariance is encouraged. We also consider this prior, but assume it is unknown and use Bayesian model selection to automatically learn the variances that control symmetry constraints from data. Further, naive residual pathways of Eq. (4) require polynomially many parameters  $C^\prime CX^4 +C^\prime CX^2$  , or  $C^\prime CX^4 +C^\prime CS^2$  for  $S\times S$  filters, using square features  $X = Y$  for simplicity. This makes using residual pathways in common architectures infeasible (see line 3 in Table 2). We overcome this by proposing factorisations and sparsifications bringing parameter counts of relaxed equivariance closer to convolutional layers.
Factored layers Inspired by common factorisations in convolutional architectures, such as depthwise [Chollet, 2017] and group separability [Knigge et al., 2022], we propose to factor fully-connected layers between input and output dimensions  $\pmb{\theta}(c', c, x', y', x, y) = \pmb{\theta}_1(c', c, x', y')\pmb{\theta}_2(c', c, x, y)$ . This greatly reduces the parameters, while fully maintaining the benefits of fully-connected structure from full spatial  $H \times W$  input and output  $H' \times W'$  to both input and outputs channels  $C \times C'$ .
$$
\sum_ {c} \sum_ {x, y} \mathbf {x} (c, x, y) \boldsymbol {\theta} \left(c ^ {\prime}, c, x ^ {\prime}, y ^ {\prime}, x, y\right) = \underbrace {\sum_ {c} \boldsymbol {\theta} _ {1} \left(c ^ {\prime} , c , x ^ {\prime} , y ^ {\prime}\right) \sum_ {x , y}} _ {\text {f a c t o r e d f u l l - c o n n e c t e d (F - F C)}} \tag {5}
$$
The factorisation of Eq. (5) has several benefits. First, it reduces the number of parameters required for the expensive the fully-connected path from  $C'CX^2Y^2$  to  $2C'CXY$ . This is essential in practice to obtain parameter counts close to convolutional layers (133.6M→1.5M in Table 2). Secondly, we can write the layer as a composition of two linear operations summing over  $\sum_{x,y}$  and  $\sum_{c}$ , which we use to derive Kronecker-factored approximations to the second-order curvature and enable the use of scalable Laplace approximations. Lastly, the parameters  $\theta_1$  and  $\theta_2$  have a single spatial domain that lends themselves to further optional sparsification with basis functions.
Figure 1: Eq. (5) visualised as tensors.
Spatial sparsification To further improve the parameter efficiency, we consider parameterising spatial dimensions in a lower dimensional basis  $\bar{s}_{\pmb{\theta}}^{\omega}(c', c, \bar{x}, \bar{y}) = \sum_{j=1}^{M} a_j \phi_j^{c', c}(\bar{x}, \bar{y})$ . Our approach is partly inspired by B-spline Lie
Figre 2: Spatial sparsification  $s^{\omega}$  Illustrated.
group filters of [Bekkers, 2019]. We follow van der Ouderaa and van der Wilk [2023] proposing standard exponential basis functions  $\phi_j^{c',c}(\bar{x},\bar{y}) = u_j\exp (-\omega^2 ([\bar{x}\bar{y} ]^T -z_j)^2)$  for convolutional filters, with anchor point locations  $z_{j}\in \mathbb{R}^{2}$  and anchor point values  $u_{j}\in \mathbb{R}$ . The work demonstrated that lengthscales  $\omega$  can allow differentiable control over equivariances (used in Section 3.2) as well as practical sparsification with negligible performance loss, even if filters are already small (e.g.  $3\times 3$ ).
$$
\sum_ {c} \sum_ {x, y} \mathbf {x} (c, x, y) \bar {\boldsymbol {\theta}} \left(c ^ {\prime}, c, x ^ {\prime} - x, y ^ {\prime} - y\right) = \underbrace {\sum_ {c} \sum_ {x , y} \mathbf {x} (c , x , y) \bar {s} ^ {\omega} \left(c ^ {\prime} , c , x ^ {\prime} - x , y ^ {\prime} - y\right)} _ {\text {s p a r s e c o n v l u t i o n (S - C O N V)}} \tag {6}
$$
This reduces the number of parameters from  $C'CS^2$  to  $C'CP$  with  $P$  anchor points, or  $C'CP + 2P$  if anchor point locations  $\{z_j\}_{j=1}^M$  (shared over channels) are also considered learnable parameters. As weights in our proposed factored F-FC have independently factored spatial input and output domains, we can perform a similar sparsification. We write  $s_1^\omega$  and  $s_2^\omega$  for spatially sparsified  $\theta_1$  and  $\theta_2$ :
$$
\sum_ {c} \boldsymbol {\theta} _ {1} \left(c ^ {\prime}, c, x ^ {\prime}, y ^ {\prime}\right) \sum_ {x, y} \mathbf {x} (c, x, y) \boldsymbol {\theta} _ {2} \left(c ^ {\prime}, c, x, y\right) = \underbrace {\sum_ {c} \boldsymbol {s} _ {1} ^ {\omega} \left(c ^ {\prime} , c , x ^ {\prime} , y ^ {\prime}\right) \sum_ {x , y} \mathbf {x} (c , x , y) \boldsymbol {s} _ {2} ^ {\omega} \left(c ^ {\prime} , c , x , y\right)} _ {\text {s p a r s e f u l l - c o n n e c t e d (S - F C)}} (7)
$$
reducing the parameter cost from  $2C' CXY$  to  $C' CP$  with  $P$  total number of anchor points. Compared to convolutional filters, which already have small spatial supports, sparsification of fully-connected layers will typically result in a much larger absolute reduction of parameters, for fixed  $P$ . As discussed in App. A, sparsified layers of Eqs. (6) and (7) can be extended to other groups, but require well-defined basis functions  $\phi$  on groups, as explored in [Azangulov et al., 2022].
3.2 Specifying the amount of equivariance in the prior
To allow for learnable equivariance, we explicitly specify symmetry constraints in the prior and then empirically learn them through approximate Bayesian model selection. We choose a Gaussian form  $\mathcal{N}(\mathbf{0},\sigma_l^2\bar{\mathbf{I}})$  and treat variances as hyperparameters  $\pmb {\eta} = [\sigma_1,\sigma_2,\dots ]$  . We carefully parameterise the network, such that the prior in each layer  $l$  enforces an equivariance symmetry in the following limit:
$$
\sigma_ {l} ^ {2} = 0 \Longrightarrow \text {s t r i c t e q u i v a r i a n c e (E q . (2))} \quad \text {a n d} \quad \sigma_ {l} ^ {2} > 0 \Longrightarrow \text {r e l a x e d e q u i v a r i a n c e} \tag {8}
$$