A Generative Model of Symmetry Transformations
Authors: James Urquhart Allingham
Bruno Kacper Mlodozeniec
Shreyas Padhy
Javier Antoran
David Krueger
Richard E. Turner
Eric Nalisnick
José Miguel Hernández-Lobato
Venue: NeurIPS 2024 poster (2024)
Keywords: approximate symmetries
invariances
deep generative models
Primary area: generative_models
TL;DR: We construct a generative model that learns which (approximate) symmetries are present in a dataset and leverage it for improved data efficiency

Abstract: Correctly capturing the symmetry transformations of data can lead to efficient models with strong generalization capabilities, though methods incorporating symmetries often require prior knowledge.
While recent advancements have been made in learning those symmetries directly from the dataset, most of this work has focused on the discriminative setting.
In this paper, we take inspiration from group theoretic ideas to construct a generative model that explicitly aims to capture the data's approximate symmetries. 
This results in a model that, given a prespecified broad set of possible symmetries, learns to what extent, if at all, those symmetries are actually present.
Our model can be seen as a generative process for data augmentation.
We provide a simple algorithm for learning our generative model and empirically demonstrate its ability to capture symmetries under affine and color transformations, in an interpretable way.
Combining our symmetry model with standard generative models results in higher marginal test-log-likelihoods and improved data efficiency.

A Generative Model of Symmetry Transformations
James Urquhart Allingham
Bruno Kacper Mlodozeniec University of Cambridge
University of Cambridge jua230cam ac.uk
MPI for Intelligent Systems, Tübingen
bkm28&cam. ac.
Shreyas Padhy University of Cambridge sp2058@cam. ac. uk
Javier Antorán
David Krueger
University of Cambridge
University of Cambridge
Ångstrom Al ja6660cam. ac uk
david. scott wtw.otttrrruegeggg....oom
Richard E. Turner University of Cambridge ret260cam. ac. uk
Eric Nalisnick University of Amsterdam e. t malinaiiii ssnickkuua.n.
José Miguel Hernández-Lobato University of Cambridge jmh2330cam ac. uk
Abstract
Correctly capturing the symmetry transformations of data can lead to efficient models with strong generalization capabilities, though methods incorporating symmetries often require prior knowledge. While recent advancements have been made in learning those symmetries directly from the dataset, most of this work has focused on the discriminative setting. In this paper, we take inspiration from group theoretic ideas to construct generative model that explicitly aims to capture the data's 's approximate symmetries. This results in a model that, given prespecified broad set of possible symmetries, learns to what extent, if all, those symmetries are actually present. Our model can be seen as generative process for data augmentation. We provide simple algorithm for learning our generative model and empirically demonstrate its ability to capture symmetries under afline and color transformations, in an interpretable way. Combining our symmetry model with standard generative models results in higher marginal dest-log-likeliben and improved data efficiency.
Introduction
An example of symmetry-aware generative process that we aim to model this pa- per. A prototype x ( is transformed by T, into an observation x (
Figure 1: Left:
X
y. The transformation of possible instances of x can result from applying Titwith few elements shown. Under this generative process, the prototype is an arbi- trary orbit element. Each element in the orbit has probability p (x x) induced by p (11 x). E.g.. for handwritten '3's, we expect digits in an upright orientation with some rotation around, say =40", corresponding to natural variations in handwriting.
3
3
3
c.g., rotation-is is parameterized by nJ-cgg, an gle. Right: The corresponding orbit-i.e., the
3
T (x)
3
p x)
3
(n ())
11
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Many physical phenomena exhibit symmetries; for example, many of the observable galaxies in the night sky share similar characteristics when accounting for their different rotations, velocities, and sizes. Hence, if we are to represent the world with generative models, they can be made more faithful and data-efficient by incorporating notions of symmetry. This has been well-understood for discriminative models for decades. Incorporating inductive biases such as invariance or equivariance to symmetry transformations dates back (at least) to ConvNets, which incorporate translation sym- metries [LeCun et al., 1989]-ann and can be extended to reflection and rotation [Cohen and Welling, 20d6]-and more recently, transformers, with permutation symmetries [Lee et al., 2019].
In many cases, it is not known priori which symmetries are present in the data. Learning symmetries in discriminative modeling is an active field of research [Nalisnick and Smyth, 2018, van der Wilk et al., 2018, Benton et al., 2020, Schwöbel et al.. 2021, van der Ouderaa and van der Wilk, 2022, Rommel et al., 2022, Romero and Lohit, 2022, Immer et al., 2022, 2023, Miao al., 2023, focus on invariant discriminative
Mlodozeniec et al., 2023]. However, in these models the label is often assumed to be invariant, and thus, the symmetry information can be removed rather than explicitly modeled. On the other hand, generative model must capture the factors of variation corresponding to the symmetry transformations of the data. Doing so can provide benefits such better representation learning-by disentangling symmetry from other latent variables [Antorán and Miguel, 2019] and data efficiency due to compactly encoding of factor(s) of variation corresponding to symmetries. Furthermore, learning about underlying symmetries in data could be used for scientific discovery.
We propose generative model that explicitly encodes the (partial) symmetries in the data. Here, we are primarily interested in using this model to inspect the distribution over naturally occurring transformations for given example x, and resample new "naturally" augmented versions of the example. Our contributions are
1.
We propose Symmetry-aware Generative Model (SGM). The SGM's latent representation is separated into an invariant component x and an equivariant component n, The latter, п, captures the symmetries in the data, while x captures none. We recover x by applying a parameterised transformation, x Tn (8). We call & prototype since each & can produce arbitrarily transformed observations; see Figure 1.
2. We propose a two-stage algorithm for learning our SGM: first learning & using a self- supervised approach and then learning n via maximum likelihood. Importantly, this does not require modeling the distribution of prototypes p (&), allowing the procedure to remain tractable even for complex data.
3. We verify experimentally that our SGM completely captures affine and color symmetries. A VAE's marginal test-log- likelihooo can improved by using our SGM to incorporate symmetries. Additionally, unlike a standard VAE, explicitly modeling symmetries makes our VAE-SGM hybrid robust deleting half of the dataset.
Notation. We use a, a, and A (i.c., lower, bold lower, and bold upper case) for scalars, vectors, and matrices, respectively. We distinguish between random variables such as n,, A, and their realizations æ, n, A. Thus, for continuous a, (a) is a PDF that returns a density (a = a) = p(a). We use to represent function composition, e.g., f2.
2 Symmetry-aware Generative Model (SGM)
Consider dataset of observations (a,, ( N on a space X' and collection (T₂) of transformations Tn : X' - X' parameterised by transformation parameters 1 E H. C Rdd. We assume (T₂) )ee (abbreviated )T,)) )) form group. Loosely, our aim is to model the distribution over transformations present in the data. To do so, we model the distribution (xx) by decomposing into two disparate parts: (1) a distribution over prototypes and (2) distribution over parameters controlling transfor- mations to be applied to a prototype. Concretely, we specify our generative model as follows (also depicted in Figure 2):
* p(x),
1 ~ Pvp (1 x)
(2)
(x).
(3)
2

---

5
s (x)
-300 n
Tn (x)
mse (:
Lsst.
p mmmd)
-500
7md
Tn (x)
Lmd
fw (x)
-800
T (x)
Figure 4: Self-supervised symmetry learning. We encourage fw (x) to be equivariant by mapping and randomly transformed z to the same è. Gray text shows examples for each variable in the graph. Note that and 2'dd may not appear in the dataset; see Figure
That is, the SGM assumes that each observation x is generated by applying transformation Tn parameterized by a latent variable 11-1o latent prototype &. Since &, by assumption, contains no information about the symmetries in the data, Pv (1 8) must model the distribution over the transformations T present in the data.
Motivation. Why would expect specifying (x) in this way to be useful? Firstly, our SGM allows us to query a distribution over naturally occurring transformations pv (n | = T, (a)) for any input æ, given the matching prototype a := T, (ar). Secondly, we expect our SGM to align with the true physical process of gener- ating the data for many interesting datasets. As an illustrative example, when person writes digit, they first decide what kind of digit to e.g. the prototype could be an upright '3'-but when they put pen to paper, the digit they pictured is transformed due to various factors governing their handwriting'. Similarly, when a photographer captures an object, the photo is also function of latent factors of variation, such as lighting, the lens, camera shake, etc.
generative
1/
inference
implicit
invariant equivariant
N
Figure 2: SGM graphical model. The implicit edges denote that & is fully specified by and x- since & = T (x)-ann and thus only n needs to be inferred given and observation
What do we require of prototype? & can informally be considered a canonical/refeerncce example with transformation applied it. More precisely, require that for any orbit of element
x defined as the set of elements in X' which x can be mapped to by a transformation in [Tn there
exactly one prototype in the orbit. Figure depicts example orbit-a set rotated variants of a -3'-with a unique prototype.
....) of all
Why do we want a group? Having the transformations (Tn} be a group simplifies things, since {Tn) will then naturally par- tition the space X' into (disjoint) orbits. Within each orbit, every element can be transformed into one another with transforma- tion in {Tn}. ). As an example of such partition, if our collection of transformations were horizontal shifts Tn x + x + (n, 0) acting point x E R22 then the different orbits will correspond to all points on given horizontal line; see Figure: 3. Therefore, if have chosen a unique prototype for each orbit and {Tn} forms group, any two elements x, x' E X' will have the same prototype and only if they can be transformed into one another.
Figure 3: Orbits due to horizon- tal shift transformations. Each point (æ1, x2) is transformed via Ty (x1, (22 ) -> (x1, (2) (22 (1. 0). Thus, horizontal lines form dis- joint orbits which any point can be transformed into any other point on the same line but not on another line. For each line, we can choose an arbitrary prototype from which all other points on the line can be reached via Ty-
In Section 2.1, we describe method for learning a transformation inference function fw X' - H. with parameters w. that for x E X returns transformation parameters E H. as = fa (x).
These map x to a prototype & : -x that generates x = Tn 2 We then apply standard generative modeling tools to learn p (x, n) p ((x) Pv (1 x) given the generated data pairs (Φ,, 7n 1
'Our SGM does not always perfectly match data gettraeenraaing process. E.g., a person is unlikely "imagine" the same prototype for both a '6' or a which can often he transformed into one another with rotation. Thh transformation is not necessarily unique.
3

---

2.1 Learning
We now discuss learning for the two NNs required by our model, fw (x) and Py (n x). In Ap- pendix A, we connect our learning algorithm with MLL optimization using an ELBO.
Transformation inference function. For T with n given by fw. to map x to a prototype x, it must, by definition, map all elements in any given orbit to the same element in that orbit. In other words, the output of (w) (a:) should be invariant to transformations Ty of x::
Fi (z)) (z) = Till (T,)(())) (z)), Vny' E H.
(4)
To learn such a function, we optimize for this property directly. To this end, we sample transformation parameters 7nd from some distribution over parameters pmmma). This allows us to get random samples Fnd := T (x) E X' in the orbit of any given element æ E X. Since we want full (i.e., strict) invariance, p maa) must have support the entire orbit [van der Ouderaa and van der Wilk, 2022]. We then learn an equivariant via a self--uuerrised learning (SSL) scheme inspired by methods like BYOL [Grill al., 2020] and, more directly, BINCE [Dubois et al., 2021]. For example, we could use the objective illustrated in Figure 4:
[ (FFma) FTulle (2c) (1o) (1=) Fmd Tim (2r), 7odd P(1ma))
Our actual objective differs slightly. Since (ay) (x=" implies xr' =Tp/ Tap (ar"),
Tr (m) TT (mmaa (Ima)
This change allows us reduce the number of small discretization errors introduced with each trans- formation application by replacing repeated transformations with a single composed transformation; see Section 3.1 for further discussion. Our SSL loss is given in line of Algorithm I.
Generative model of transformations. Once we have a prototype inference fune- tion, we simply learn Psb (n x) by maxi- mum likelihood
Algorithm Learning
Require:
initial parameters Winn &
Vinn dataset
1:
function SSL SSL_LOSS(x;44
the created data pairs (a in line Algorithm While we need to specify the kinds of symmetry transfor- mations Tn we expect to see the data, by learning Pv (n x) the model can learn the degree to which those transformations present in the data. Thus,
2: nu f fw (x))
(fw Tiul (a.) (z:)). This is shown
3:
Tend p (Trma)
4:
Bmd <- T (x)
nans Sw (Frdd)
x' + T T (Irnd)
7:
output mse( x'
function MLE LOSS(W w, 1))
spec- ify several potential symmetry transfor- mations and learn that some are absent in the data. Furthermore, the required prior knowledge (the support of p mmm)) is small compared to what our SGM can learn (the shapes of the distributions for each of the present transformations).
9:
ni + (w)
10:
à -- Tx-
11: output - -log (N (:)
12: w, vb f- Winit: Viiiit
13: while not converged do
14: X+ next_ (atch(D
15: update w with V. 1 Σ
SSL LSSS
16: while s/s not converged do
Since are primarily interested in using the model to (a) inspect the distribution
17:
X +- next mext_batch(()
18: update &b with V 1 E,
X,, w. w))
naturally occurring transformations
19: output w, 2b
B
for a given element æ, and (b) resample new "naturally" augmented versions of the element, we do not need to learn (x)). We do (a)
by querying p (1 x = 2)) for à := (=)), (z), and we can do (b) by sampling 17 ~ p (n ()) and Tr, (æ). Of course, if one wanted to sample new prototypes, one could
transforming the & to get fit Pe (x) using, e.g., a VAE. Not learning p (X) greatly simplifies training for complicated datasets that would otherwise require large generative model, observation made by Dubois et al. [2021].
'I fas is equivariant by construction, our SSL scheme is unnecessary. Alas, such constructions are unknown for many transformation,, like those in this paper. Thus, provide general method for learning equivariances.

---

8
(a) Distribution for given x and X.
(m x, Φ)
.n/8 8
./
8 8 0 - 09) + 0.5 0( 1809)
8 8 0,5- 8(n - 309) + 0.5 0( + 150%)
8
8
0.5 8(n + 30") 0.5 0(n -500)
8 8
8
(b) Simple pv (n (x)
(c) Flexible py (n x)
Figure 5: Idealised examples of simple and flexible learned distributions over angles Pup (nl ()
-given the true distribution (n x) _x P (1 x, x)- 444
3 Practical Considerations and Further Motivations
Training our SGM, while simple, has potential pitfalls in practice. We discuss the key considerations in Section 3.1 and provide further recommendations in Appendix B. We then provide motivation for several of our modeling choices in Section 3.2.
3.1 Practical Considerations
Working with transformations. introduce unwanted artifacts such as blurring. For many useful transformations, we can compose transformations before applying them. For affine transformations of images, for example, we can directly multiply
Repeated application of transformations-ee.. in Figure c-can
matrices. More generally, if there is some representation of the transformation parameters T(n where composition can be performed-e.g.
matrix multiplication Tho o Tm TTo )'())' in the case where T' is group representation-ttenn we recommend composing transformations in that space to minimize the number of applications.
Partial invertibility. In many common settings, transformations are not fully invertible. We encounter two such issues when working with affine transformations of images living in finite, discrete coordinate space. Firstly, affine transformations are only approximately invertible in the discrete space due to the information loss when interpolating the transformed image onto a discrete grid. Thus, while only a single prototype x exists for any x, it may not be clear what the correct prototype is. Secondly, transformations can cause information loss due to the finite coordinate space (e.g.. by shifting the contents of the image but-od--ounds"). If appropriate bounds are known a priori, we can prevent severe information loss by constraining Mmin and Mmax using tanh, scale, and shift bijectors. Alternatively, we can augment the SSL. loss Algorithm with an invertibility loss
Linvertiiilityy ( (w) = mse ax, (i)( (T (a:i))-
Learning pap (1) (8) with imperfect inference. In practice, our transformation inference network fai (x) will not be perfect; see Figure 10. Even after training, there may be small variations in the prototypes x corresponding to different elements in the orbit of x. To make pv (1xx x) robust to these variations, we train it with prototypes corresponding to randomly transformed training data points. L.c. we modify the MLE objective in Algorithm as log (7 .,)), where a' = =T (Taye (z)) (T (x)) as in our SSL objective. Averaging the loss over multiple samples e.g., 5- Mnd is beneficial.
3.2 Modelling Choices
We now motivate some of the design choices for our SGM by means of illustrative examples. each case, assume that Tr is counter-clockwise rotation; thus, n is the angle.
1. The distribution Ps (nl {) implemented as a normalising flow. Consider a dataset of '8's rotated in the range -0 to 30:: (8, (8...., 8, .... 8) J. Let us assume that the prototype is '8'. Figure 5a shows p (n x, x), an example of the true distribution for n given x and &, for several observations, under the data generating process These distributions are composed of deltas because
"This can occur in practice since our SSL objective- which aims to make prototypes as similar as possible- trivially be minimized by removing all of the contents of an image.
'Because '8' is symmetric, p (n x, x) could be any convex combination of the two delta distributions. However, for more realistic example, consider a prototype '8' with a smaller upper loop. In this case, the p (n| &) must be bimodal to capture '8's with both smaller upper and lower loops.
5

---

(a) Distribution for 1l given x and x.
x:
x
(n x; a))
2
2
ö(n - 09)
2
2
5(n 30))
n
n|2
_n8
2-
2
+ 5
8
8
0.5
0.5 - -0) 0.5
180%)
8
8
+111+1500 150
8 8
0.5
5(n + 30") 0.5 150%)
(b)
(c) n x
Figure 6: Examples of learned distributions over angles Ptp (+) 2 with and without dependence
on X, given the true distribution p (-)-
only certain values of n will transform & into Figures 5b and 5e compare idealised examples of the learned Pq. (n giveiven& a simple uni-modal Gaussian family and a more flexible bi- bo-ma mixture- of-Gaussian family-with the aggregate true distribution (n x) -
8.....8..... P (n X, x). Here, the simple uni-modal distribution is clearly worse than the bi-modal distribution due to the large amount of probability mass being wasted on angles with low density under the true data- lata-generating process. Of course, one might argue that the bi-modal distribution is also not flexible enough. Furthermore, 'flexible enough' will be problem-specific. We solve this problem with normalizing flows, which can match a wide range of distributions.
2. The transformation parameters n depend on the prototype x. Consider a dataset of '2's and '8's rotated in the range -30 to 30": (2,.... 2, ........ 2, 8,.... 8..... }. with prototypes '2' and '8', Figure 6a shows p (1 x, x). an example of a true distribution over n, for several observations. Figures 6b and 6c compares idealised examples of learned distributions over n and n xt. Without dependence on &, the model must place probability mass between 150 and 150", in order to capture the symmetries of the '8's, however this results invalid digits digits-such as 12, è, which do not come from true data distribution. On the other hand, when depends on &, the distribution conditioned on the prototype for the '2's only needs to place mass in [-30P, 30%)
3. The prototype x is fully invariant to trans- formations of x. Models such as CNNs are most useful when we know a priori which sym- metries are present in the data. However, in many cases, this must be learned. In the case of handwritten digit recognition, we know that the model should be invariant to some amount of rotation since people naturally write with some variation in angle. But model that is invari- ant to rotations in the full range [-180%, 180"] might be unable distinguish between '6" and '9". Thus, in the literature for learning invari- ances in the discriminative setting, it is common to learn partially invariant functions that capture some degree of invariance [van der Wilk al., 2018, Benton et al., 2020, van der Ouderaa and van der Wilk, 2022]. However, as we will now show, this approach is unsuitable for our SGM, as it breaks our assumption that & contains no information about the symmetries in the data.
(a) p() x, x) with different levels of invariance.
(a) FULL
(b) PARTIAL
(c) NONE
w á p
w. @)
Φ
(n a a)
p(n x., a)
2 2 8(n - 0")
2 8(n 15")
8(n 0²)
2 2 8(n - 30")
2 8(n 15")
2 8(n - 0%)
2 2 8(n + 3
2
8(n 0*)
2 8(n 0%)
n|2
n
n n|2
n/
(b) FULL
(c) PARTIAL
(d) NONE
Figure 7: Examples of learned distributions over
angles 74p (n x) 8)- x + -with different de-
grees of invariance in the prototype x, given the true p (n {)- 1144
Consider a dataset of '2's rotated in the range -30" to 30%: predicted prototypes and the corresponding distributions over for several observations. There are three cases: (a) fully- hully-invariant& x. i.c., there is single prototype. (b) a partially-invariant x, for which there are two prototypes in this example, and (c) a non-invariant X, which takes the partially- invariant case to the extreme and has as many prototypes as observations. In the partially-invarian and non- invaiiariann cases, we can get multiple prototypes rather than single unique prototype per orbit, which is invalid under the generative model of the data. As a result, Pv (n (x) does not represent the distribution of naturally occurring transformations of & in the data. This is illustrated in Figures 7b to 7d, which show idealized examples of the learned pv (n x) each case. While the distribution in Figure 7b matches the distribution of transformations in the dataset, in Figures 7c and 7d we see that the distributions corresponding to non-unique prototype do not. To illustrate why this problem, let say would like to probe the probability of a particular transformed variant
2). Figure 7a shows

---

(a) dSprites under affine transformations
/
2
2
3
3
4
5
6
6
7
7
Z
8
8
7
9
3
3
4
6
5
6
6
4
2
8
8
6
6
0
I
2
2
3
3
4
4
6
9
7
7
2
8
8
1
6
(b) MNIST under affine transformations
0
O
1
/
2
2
3
3
4
4
5
5
6
6
7
7
7
8
8
9
9
0
O
/
/
2
2
3
3
4
4
5
5
6
6
7
7
7
8
8
9
9
0
O
/
/
2
2
3
3
4
4
5
5
6
7
7
7
8
8
9
9
(c) MNIST under color transformations
(d) GalaxyMNIST under affine and color transformations
Figure 8: Top: samples from the test set. Mid: prototypes for each test example. Bot: resampled versions of each test example given the prototype. Prototypes for examples from the same orbit (and in some cases from distinct but similar orbits) match (e.g., their size, position, rotation, etc. are similar). Resampled examples are usually indistinguishable from test examples.
of observed example. For example, given example of digit '3', want to know the probability of observing W that digit rotated by -90". Assuming we can find a prototype we would like (n & .) represent all naturally occurring augmentations. Unless x is unique, this won't 't necessarily be the case, as illustrated in Figure 7.
4 Experiments
In Section 4.1, explore SGM's ability learn symmetries. We show that it produces valid prototypes, and generates plausible samples from the data distribution, given those prototypes. Then, in Section 4.2, leverage SGM to improve data efficiency in deep generative models.
conduct experiments using three datasets- dSprits [Matthey et al., 2017]. MNIST, and GalaxyM- NIST (Walmsley et al., 2022]-and two kinds of cansformations--rffee when working with MNIST under affine transformations, we add a small amount of rotation (in the range[-159 152)) to the original data make rotations in the figures easier to For MNIST under color transformations, we first convert the grey-scale images to color images using only the red channel. We then add random hue rotation in the range [0, 0.67] and random saturation multiplier in the range (0.6, 0.9]. In the case of dSprites, we carefully control the rotations, positions, and sizes of all of the sprites. For example, in the and set the y-positions to be bimodal in the top and bottom of the images. Further details about the dSprites setup, as well all other experimental details, can be found in Appendix C. We focus on learning affine transformations (shifting, rotation, and scaling) as they are expressive while still being group that is easy work with. We also learn color transformations (hue, saturation, and value). See Appendix C.7 for details about how we parameterize Tn in both cases.
and color. In Section 4.
of the heart sprites,
have removed the rotations
4.1 Learning Symmetries
Exploring transformations and prototypes. Figure 8 shows that for both datasets and kinds of transformations we consider, our SGM produces close-to-invariant prototypes as well as re- alistic "natural" examples that almost indistinguishable from test examples. There are sev-
7

---

eral illustrative examples which bear further discussion. The heart sprites in Figure 8a show that our SGM was able learn the absence of transformation (namely rotation) in the dataset. As expected, all of the prototypes for the sprites of the same shape are the same, since these shapes are the same orbit as one another. This behaviour is also demonstrated for MNIST digits in Figures 19 and 20, The '6', '8', and '9' digits in Figure 8b demonstrate the ability of our SGM to learn bimodal
2
2
9
distributions (on rotation this case). The figure's 's third '7' is interesting because our SGM interprets it as '2'.
9
Figure From left to right, test examples, their prototypes, and the corresponding marginal distri- butions Py, (7, xx) over translation in ær translation in y, rotation, seal- ing in x, and scaling in y.
Flexibility is important. In 17, each dimension corresponds to a different transformation. We refer to Pv (11 x) as the marginal distribution of a single transformation parameter. Fig- ure 9 shows these marginal learnt distributions for several digits from Figure 8b. We see that each of the parameters has its own range and shapes. For rotations, which are easy to reason about, we see distributions that make sense theseetheeee round "0" has an almost uniform distribution over rotations, and the '1' and one of the '9's are strongly bimodal as expected. The other '9', which does not look as much like an upside-down '6', has much smaller 2nd mode. The '2', which looks somewhat like an upside-down '7', is also bimodal. We see that prototypes of different sizes result in corresponding distributions over scaling parameters with different ranges. Figure 21 provides additional examples for MNIST with affine transformations, while Figure 22 provides the same for color transformations, and Figure 23 investigates the distributions for dSprites. These results provide experimental evidence of the need for flexibility in the generative model for Py (17 x), as conjectured in Section 3.2. We also find significant dependencies between dimensions of 1) (e.g.. rotation and translation in dSprites).
Invariance of fw and the prototypes. In Figure 10, we investigate the imperfections of the inference network by considering an iterative procedure in which prototypes are treated as observed examples, allowing us to infer a chain of successive prototypes. We show several examples of such chains, as well as the average magnitude of the trans- formation parameters at each iteration, normalized by the maximum magnitude (at iteration 0). The first prototype x1 most different from the previous &o = with successive prototypes being similar visually and as measured by the magnitude of the inferred transformation parameters. How- ever, the magnitude of the inferred parameters does not tend towards O, rather plateauing around 5% of the maximum. This highlights that, although simple NNs can learn to be approximately invariant, natively invariant architecture has the potential to improve performance.
1.0
Normalsed Avomge (Iny la
0.8
2
2
2
2
0.6
9
6
6
6
0.4
8
g
8
8
9
6
6
6
8,
82
IIoeracionss
Figure 10: Iterative prototype infer- ence. Left: starting with a test ex- ample x, we get prototype x1, then treating prototype x; as an observed example we predict the next prototype 8 Right: The average magnitude of the transformation parameters as function of iterations of this process.
4.2 VAE Data Efficiency
We use SGM to build data-efficient and robust generative models. In Figure 11, we compare a standard VAE to two VAE-SGM hybrid Modells-"AuVAAe and "Inv VInyVAE--fo different amounts of training data and added rotation of the MNIST digits. When adding rotation, each x in the dataset set is always rotated by the same angle (sampled uniformly between 0 the maximum added rotation angle). Thus, adding rotation here is not data augmentation. VAg a VAE that uses SGM re-sample transformed examples x' Tn æ (æ), introducing data augmentation at training time. VAE is a VAE that uses our SGM to convert each example æ to its prototype e at both train and test time. That is, the VAE in Inv sees only the invariant representation of each example. We also compare against VAE trained with standard data augmentation". We use test-set importance- weighted lower bound (IWLB) [Domke and Sheldon, 2018] of p (x), estimated with 300 samples of the VAE's s latent variable z, and for Inv VAE. to compare the models. Reconstruction error provided in Appendix E. Further ettails-eg... hyperparameter sweeps-art are in Appendix C.
We use rotation 41115 15"), zoom u(-10%, 10%). and x/y- u[xx2x, 2px).
8

---

As expected, for VAE as we decrease the amount of training data (- randomly added rotation, performance degrades. This is because the VAE sees fewer training examples per-degree of rotation. On the other hand, the AugVAB (Foee) is more data efficient. Its performance is unaffected by reducing the number of observations by three quarters. Furthermore, while the performance Aug and the standard VAE almost identical for small angles and large training sets, the drop in performance of AugVAE for larger random rotations is significantly smaller, Aug VAg does not see less training examples per-degree of rotation. VAE (1rr)), which natively incorporates the inductive biases of our SGM and obtains a 500 nat larger likelihood than the other models. Its performance is almost perfectly robust to rotation in the dataset. Additionally, its metrics barely change (< 10%) when trained on half the data. Finally, while the VAE with data augmentation (ree)) improves on the standard VAE for less training data, is substantially worse in the presence of more data. This contrasts our VAEs, which are almost always better. This poor performance is because the augmentations are independent of the samples. Thus, highly rotated digits can be rotated too much, smaller digits become too small, and digits near the image edges are moved out of frame. This highlights the importance of augmenting data accordance with the true data distribution.
1800
or increase the amount of
1700
1300
1200
IWLB
1100
1000
VAB
12500
AugVAF
25000
InVVAE
37500
VAE
50000
Maximum Added Rotation Angle (*)
Figure 11: Incorporating symmetries improves data efficiency. Importance- weighted lower bound (IWLB) (mean and std. eTT. over 3 random seeds) on rotated MNIST for standard VAE (w. and w.o. data aug.) and two VAE variants that incorporate symme- tries via our SGM. Improved data effi ciency demonstrated by better perfor- mance with less training data and less sensitivity to added rotation.
We further validate these results with the more complex GalaxyM- NIST dataset and enlarged set of both affine and color transfor- mations. As with our rotated MNIST with affine transformation results, in Figure 12, we see that Aug VAE VAE standard VAE to training with only half of the dataset. Our SGM captures the true data distribution with only 3500 training examples.
12030
10030
outperforms the Furthermore, we see that Aug is robust
6000
Aug/All
7000
Num. Thain
5 Related Work
GalaxyMNIST data- efficiency (3 seed mean & std. err.).
Figure 12:
Learning Lie groups. Rao and Ruderman [1998], Miao and Rao [2007]. Keurti et al. [2023] learn Lie groups from sequences of transformed images in unsupervised fashion. Hashimoto al. [2017] learn to represent an image as a linear combination of transformed versions of its nearest neighbors. Dehmamy al. [2021] use Lie algebras define CNNs for automatic symmetry discovery. Yang et al. [2023] use a GAN-based approach to learn transformations of examples that leave the original data distribution unchanged, thereby fooling a discriminator. Falorsi et al. [2019] introduce a reparameterization trick for learning densities on arbitrary. but known, Lie groups Chau et al. [2022] learn a generative model over Lie group transformations applied to prototypical images that are themselves composed of sparse combinations of learned dictionary elements.
Learning a prototype. Kaba al. [2023] note that symmetry-based NNs are often contained in their architectures. Like us, they propose to learn "canonicalization functions" that produce prototypical representations of the data. Mondal et al. [2023] show that such canonicalization functions can be used to make symmetry priors, do not degrade performance. Similarly, Kim et al. [2023] learn cethitecturr-ggnnosee equivariant functions by averaging a non-equivariant function over a probabilistic prototypice input. Finally, while not explicitly trained to produce prototypes, spatial transformers learn to undo transformations such as translation, scaling, and rotations [Jaderberg et al., 2015].
NNs equivariant and, when combined with dataset- dependenn
Data augmentations and symmetries. Prior work makes several connections between data aug- mentation and symmetries relevant to our findings. Bouchacourt et al. [2021b] show that invariances in the model tend to result from natural variations in the data rather than data augmentation or model architecture. This supports approach of learning data augmentation from the data and our

---

Sabhitectureagnnostie self-supervised invariance learning method. Balestriero et al. [2022]. Miao et al. [2023]. Bouchacourt et al. [2021b] show that learned symmetries (i.e., data augmentation) should be class-dependent, much like our transformations are prototype-dependen.
Symmetry-aware latent spaces. Encoding symmetries in latent space is well-studied. Higgins et al. [2018] posit that symmetry transformations that leave some parts the world invariant are responsible for exploitable structure in any dataset. Thus, agents benefit from disentangled representations that separate out these transformations. Winter et al. [2022] split the latent space of an auto-encoder into invariant and equivariant partitions. However, they rely on geometric NN architectures, contrasting with our self-supervised learning approach. Furthermore, they do not learn a generative theel-they reconstruct the input exactly-thus, they cannot sample new observations given a prototype. Xu et al. [2021] propose group equivariant subsampling layers that allow them to construct autoencoders with equivariant representations. Shu et al. [2018] 8] propose an autoencoder whose representations are split such that the reconstruction of an observation is decomposed into a "template" (much like our prototypes) and a spatial deformation (transformation..
In the generative setting, Louizos et al. [2016] 6] construct a VAE with a latent space that is invariant to pre-specified sensitive attributes of the data. However, these sensitive attributes are observed rather than learned. Similarly, Alice et al. [2023] construct a VAE with a partitioned latent space with a component that is invariance spurious factors of variation in the data. Bouchacourt et al. [2018], Hosoya [2019] learn VAE with two latent spaces-a per-observation equivariant latent and an invariant latent shared across grouped examples. Other works have constructed rotation equivariant [Kuzina et al., 2022] and partitioned equivariant and invariant [Vadgama et al., 2022] latent spaces. Antorán and Miguel [2019], Ilse et al. [2020] split the latent space of VAE into domain, class, and residual variation components. The first of which can capture rotation symmetry in wwitteritten digits. Unlike us, they require class labels and auxiliary classifiers. Keller and Welling [2021] construct a VAE with a topographically organised latent space such that an approximate equivariance is learned from sequences of observations. In contrast the works above, Bouchacourt et al. [2021a] argue that learning symmetries should not be achieved via a partitioned latent space but rather learning equivariant operators that are applied to the whole latent space. Finally, while Nalisnick and Smyth [2017] do not learn symmetries, their information lower bound objective is reminiscent of several works above-and our own, see Appendix A-in minimizing the mutual information between two quantities when learning prior.
Self-supervised Equivariant Learning [Dangovski et al., 2022] generalize standard invariant SSL methods to produce representations that can be either insensitive (invariant) or sensitive (equivariant) to transformations in the data. Similarly, Eastwood et al. [2023] use a self-supervised learning approach disentangle sources of variation in a dataset, thereby learning a representation that is equivariant to each of the sources while invariant to all others.
Conclusion
We have presented a Symmetry-awart Generative Model (SGM) and demonstrated that it is able to learn, in an unsupervised manner, distribution over symmetries present in a dataset. This is done by modeling the observations as a random transformation of an invariant latent prototype. This the first such model we are aware of. Building generative models that incorporate this understanding of symmetries significantly improves log-likelihoods and data sparsity robustness. This is exciting the context of modern generative models, which are close to exhausting all of the data on the internet. We are also excited about the use of SGM for scientific discovery, given that the framework is ideal for probing for naturally occurring symmetries present in systems. For example, we could apply SGM to marginalize out the idiosyncrasies of different measuring equipment and observation geometry in radio astronomy data. Additionally, given the success of using our SGM for data augmentation when training VAES, it would be interesting to apply it to data augmentation in discriminative settings and compare it with methods such as Benton et al. [2020], Miao al. [2023].
The main limitation of our SGM is that it requires specifying the super-set of possible symmetries. Future work might relax this requirement or explore how robust our SGM is to even larger sets. Furthermore, care must sometimes taken when specifying the set of symmetries. For example, when rotating to images with "content" up to the boundaries of the image; see Appendix E.2.
10

---

Acknowledgements
The authors would like to thank Taliesin Beynon for helpful discussions and Emile Mathieu for providing feedback on the paper. This work has been performed using resources provided by the Cambridge Tier-2 system operated by the University of Cambridge Research Computing Service (htt://www.hhpc.cag www.wp.c.c.ama..kk... uckuk funded by EPSRC Tier-2 capital grant EP/T022159/1. This work was also supported with Cloud TPUs from Google's TPU Research Cloud (TRC). JUA acknowledges funding from the EPSRC, the Michael E. Fisher Studentship in Machine Learning, and the Qualcomm Innovation Fellowship. JUA was also supported by an ELLIS mobility grant. SP acknowledges support from the Harding Distinguished Postgraduate Scholars Programme Leverage Scheme. JA acknowledges support from Microsoft Research, through its PhD Scholarship Programme, and from the EPSRC. JMH acknowledges support from a Turing AI Fellowship under grant EP/V023756/1 RET is supported by Google, Amazon, ARM, Improbable, EPSRC grant EP/T005386/1, and the EPSRC Probabilistic AI Hub (ProbAI, EP/Y028783/1).
References
Hananch Alice, Ferdinand Kapl, Soroor Hediyeh-Zadeer and Fabian J. Theis. Conditionally invariant representation learning for disentangling cellular heterogeneity. CoRR, abs/2307.000555 2023. doi: 10.48550//rXiv.2300 000558. (Cited on p. 10.)
James Urquhart Allingham, Javier Antoran, Shreyas Padhy, Eric Nalisnick, and José Miguel HernándezzLobato Learning generative models with invariance to symmetries. NeurIPS 2022 Workshop on Symmetry and Geomelry in Neural Representations, 2022. (Cited on p. 17.)
Javier Antorán and Antonio Miguel. Disentangling and learning robust representations with natural clustering. In M. Arif Wani, Taghi M. Khoshgoftaar, Dingding Wang, Huanjing Wang, and Naeem Seliya, editors, 18th IEEE International Conference On Machine Learning And Applications, ICMLA 2019, Boca Raton, FL, USA, December 16-19, 2019, pages 694-699, IEEE, 2019. doi: 10. 1.11009CMM...200199
URL http:: org/ 10. M.110/IIMMLL 2019. .0125.. (Cited
on pp. 2 and 10.)
Randall Balestriero. Léon Bottou, and Yann LeCun. The effects of regularization and data augmenta- tion are class dependent. In NeurIPS, 2022. (Cited on p. 10.)
Gregory W. Benton, Marc Finzi, Pavel Izmailov, and Andrew Gordon Wilson. Learning invariances in neural networks from training data. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. (Cited on pp. 2, 6, 10, 17, and 24.)
Diane Bouchacourt, Ryota Tomioka, and Sehastian Nowozin. Multi-level variational autoencoder: Learning disentangled representations from grouped observations. In Proceedings of the Thirty- Second AAAI Conference on Artificial Intelligence, 2018. (Cited on p. 10.)
Diane Bouchacourt, Mark Ibrahim, and Stéphane Deny. Addressing the topological defects of disentanglement via distributed operators. CoRR, abs/2102.05623, 2021a. (Cited on p. 10.)
Diane Bouchacourt, Mark Ibrahim, and Ari S. Morcos. Grounding inductive biases in natural images: invariance stems from variations in data. In Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pages 19566-19579, 2021b. (Cited on pp. and 10.)
Ho Yin Chau, Frank Qiu, Yubei Chen, and Bruno A. Olshausen. Disentangling images with lie group transformations and sparse coding. In Sophia Sanborn, Christian Shewmake, Simone Azeglio, Arianna Di Bernardo, and Nina Miolane, editors, NeurIPS Workshop on Symmetry and Geometry in Neural Representations, 03 December 2022, New Orleans, Lousiana, USA, volume 197 of Proceedings of Machine Learning Research, pages 22-47. PMLR, 2022. URL https://proceedis
mlr. press/v197 v / chau23a. html. (Cited on p. 9.)
Taco Cohen and Max Welling, Group equivariant convolutional networks. In Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York City, NY, USA, June 19-24, 2016, volume 48 of JMLR Workshop and Conference Proceedings, pages 2990-2999. JML.Rorrg 2016. (Cited p. 2.)
11

---

Rumen Dangovski, Li Jing, Charlotte Loh, Seungwook Han, Akash Srivastava, Brian Cheung, Pulkit Agrawal, and Marin Soljacic. Equivariant self-supervised learning: Encouraging equivariance in representations. In The Tenth International Conference on Learning Representations, ICIR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL hattps://openreviex (Cited on 10.)
Nima Dehmamy, Robin Walters, Yanchen Liu, Dashun Wang, and Rose Yu. Automatic symmetry discovery with lie algebra convolutional network. In Marc' Aurelio Ranzato, Alina Beygelzimer. Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, edi- tors, Advances in Neural Information Processing Systems 34: Annual Conference on Neu- ral Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, 2021. URL https: /ttts://pproceedinggs neuripss cc/ httpppppp//////eeeee Abstrat html. (Cited on p. 9.)
pages
Justin Domke and Daniel Sheldon. Importance weighting and variational inference. CoRR, 0bss/8088090344, 2018. URL http://arxiv..rrt 1ttps///88800000 09034 (Cited on p. 8.)
Yann Dubois, Benjamin Bloem-Reddy, Karen Ullrich, and Chris Maddison. Lossy compression for lossless prediction. In Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pages
2021. (Cited on pp. 4 and 17.)
Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios. Neural spline flows. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 7509-7520, 2019. (Cited on p. 20.)
Cian Eastwood, Julius von Kügelgen, Linus Ericsson, Diane Bouchacourt, Pascal Vincent, Bernhard Schölkopf, and Mark Ibrahim. Self-supervised disentanglement by leveraging structure in data augmentations. CoRR, 00o//311.088111 2023, doi: 10. 10.48500/ARX1..31..0888 //doi.org/10 wwrgg/..885500.. 48550/ arXiv. 2311 2311.08815. 08815. (Cited on p. 10.)
URL https:
Luca Falorsi, Pim de Haan, Tim R. Davidson, and Patrick Forré. Reparameterizing distributions on lie groups. In Kamalika Chaudhuri and Masashi Sugiyama, editors, The 22nd International Conference on Artificial Intelligence and Statistics, AISTATS 2019, 16-18 April 2019, Naha, Okinawa, Japan, volume 89 of Proceedings of Machine Learning Research, pages 3244-322 3253.. PMLR, 2019. URL http: httt:://proceedingg .ww.inngm........mmm..m.....m..m wtt.:s.mm.ppprrer.e. frress/v89/falorsii 19a html. (Cited on p. 9.
Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H. Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Ávila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Koray Kavukcuoglu, Rémi Munos, and Michal Valko. Bootstrap your own latent A new approach self-supervised learning. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NewrIPS 2020, December 6-12, 2020, virtual, 2020. (Cited on p. 4.)
Tatsunori B. Hashimoto, Percy Liang, and John C. Duchi. Unsupervised transformation learning via convex relaxations. In Advances in Neural Information Processing Systems 30, 2017. (Cited on p. 9.)
Irina Higgins, David Amos, David Pfau, Sébastien Racanière, Loic Matthey, Danilo J. Rezende, and Alexander Lerchner. Towards a definition of disentangled representations. CoRR, abs/1812.02230. 2018. URL http://arxiv org/ abs/ 1812. 02230. (Cited on p. 10.)
Haruo Hosoya. Group-based learning of disentangled representations with generalizability for novel contents. In Proceedings of the Twenty-Eighth International Joint Conference on Artificial Intelligence, IJCAI, 2019. (Cited p. 10.)
Maximilian Ilse, Jakub M. Tomezak, Christos Louizos, and Max Welling. DIVA: domain invariant variational autoencoders. In International Conference Medical Imaging with Deep Learning, MIDI. 2020, 6-8 July 2020, Montréal, QC, Canada, volume 121 of Proceedings of Machine Learning Research, pages 322-348 PMLR, 2020. (Cited on p. 10.)
12

---

Alexander Immer, Tycho F. A. van der Ouderaa, Vincent Fortuin, Gunnar Rätsch, and Mark van der Wilk. Invariance learning in deep neural networks with differentiable laplace approximations. CoRR, abs/2202. 1 10638, 2022. (Cited on pp. 2 and 17.)
Alexander Immer, Tycho F. A. van der Ouderaa, Mark van der Wilk, Gunnar Rätsch, and Bernhard Schölkopf. Stochastic marginal likelihood gradients using neural tangent kernels. In International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pages 14333-14352, PMLR, 2023. (Cited on p. 2.)
Max Jaderberg, Karen Simonyan, Andrew Zisserman, and Koray Kavukcuoglu. Spatial transformer networks. In Advances in Neural Information Processing Systems 28: Annual Conference on Neural Information Processing Systems 2015, December 7-12, 2015, Montreal, Quebec, Canada, pages 2017-2025, 2015. (Cited on pp. 9 and 24.)
Sékou-Oumar Kaba, Amab Kumar Mondal, Yan Zhang, Yoshua Bengio, and Siamak Ravanbakhsh. Equivariance with learned canonicalization functions. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt. Sivan Sabato, and Jonathan Scarlett, editors, International Conference Machine Leurning, ICML 2023, 23-29 July 2023, Homolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pages 15546-15566. PMLR, 2023. URL http:::/pprceeeiingg wttpaanss.......ee//0///00
mlr. www.ppressssos /eeess/vv02// https//222//abbaab3ee kaba23a.
html. (Cited on p. 9.)
T. Anderson Keller and Max Welling. Topographic vaes learn equivariant capsules. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neurul Information Processing Systems 34: Annual Conference on Neu- ral Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pages 28585-28597, 2021. URL https: / /ptp/eeeeocegss..e//eeee/c/...se neurips. cc/ htc/pappr/222211 /22ppr/222///hh// 1097006651102200000fff1177771333
Abstract html. (Cited on p. 10.)
Hamza Keurti, Hsiao-Ru Pan, Michel Besserve, Benjamin F. Grewe, and Bernhard Schölkopf. Homomorphism autoencoder learning group structured representations from observed transitions. 2023. (Cited on p. 9.)
Jinwoo Kim, Dat Nguyen, Ayhan Suleymanzade, Hyeokjun An, and Seunghoon Hong. Learn- ing probabilistic symmetrization for architecture agnostic equivariance. In Alice Oh, Tris- tan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Advances in Neural Information Processing Systems 36: Annual Conference on Neural In- formation Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16. 2023, 2023. URL http://paperv/ nips wwwp.icppp........eeeea hetp/pppe_ffffie files wallespappn http//appppp//// /2023/ 22023/hahh/ 33677786667007777700000077776773 -ww.ess-tt-tttataeeeeeeeeeaaaaet Conttre-ooeeeeeeeeeeoe html. (Cited on p. 9.)
Anna Kuzina, Kumar Pratik, Fabio Valerio Massoli, and Arash Behboodi. Equivariant priors for compressed sensing with unknown orientation. In International Conference Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pages 11753-11777 PMLR, 2022. (Cited on p. 10.)
Yann LeCun, Bernhard E. Boser, John S. Denker, Donnie Henderson, Richard E. Howard, Wayne E. Hubbard, and Lawrence D. Jackel. Backpropagation applied to handwritten zip code recognition. 1989. 1 1..5411. (Cited on 2.)
Neural Comput., 144):541-551, 1989, doi:
Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online]. Available: http://yann. lecun. wwt..oo///xxb/mmmis...mmmm 2, 2010. (Cited on p. 20.)
Juho Lee, Yoonho Lee, Jungtack Kim, Adam R. Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for attention-based permutation-invarian neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, ICMI. 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pages 3744-3753. PMLR, 2019. URL nttp://proceedings mlr, www.m...reeeteeee www.:preesvvvv. /ewe/9de..199 html. (Cited p. 2.)
Christos Louizos, Kevin Swersky, Yujia Li, Max Welling, and Richard S. Zemel. The variational fair autoencoder. In 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings, 2016. (Cited on p. 10.)
13

---

Kaitlin Maile, Dennis George Wilson, and Patrick Forré. Equivariance architectural op- timization of neural networks. The Eleventh International Conference on Learning Rep- resentations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL (Cited on p. 17.)
https://openrevier net.ppdfiddaabrCafAA3.
Loic Matthey, Irina Higgins, Demis Hassabis, and Alexander Lerchner. dsprites: Disentanglement 2017. (Cited on pp. 7 and 21.)
testing sprites dataset. http:/ttwwh.oooeeemei/d//ppiipet
Ning Miao, Tom Rainforth, Emile Mathieu, Yann Dubois, Yee Whye Teh, Adam Foster, and Hyunjik Kim. Learning instance-specific augmentations by capturing local invariances. In International Conference on Machine Learning, ICML. 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pages 24720- PMLR, 2023. (Cited on pp. 2 and 10.)
Xu Miao and Rajesh P. N. Rao. Learning the lie groups of visual invariance. Neural Computation, 19 (10):2665-2693, 2007. (Cited on p. 9.)
Bruno Kacper Mlodozeniec, Matthias Reisser, and Christos Louizos. Hyperparameter optimization through neural network partitioning. In The Eleventh International Conference on Learning Representations, 2023. (Cited on p. 2.)
Arnab Kumar Mondal, Siba Smarak Panigrahi, Oumar Kaba, Sai Mudumba, and Siamak Ravanbakhsh. Equivariant adaptation of large pretrained models. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Ad- vances in Neural Information Processing Systems 36: Annual Conference on Neural In- formation Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December htpppaap////// pppppr//23///// 22////222///ha/// 10 16, 2023, 2023. URL http://papers.aet wippppppprs.ps...ppppppppppppppm cc/ htpppppppllleeeeeeee wtt.fffaapeeee/////eeeeeaaaaaeee 94BB6618002e1333355500..o22
-wte:e222--bstaracc Abstract Conf Conference....oo. httml. (Cited on p. 9.)
Eric T. Nalisnick and Padhraic Smyth. Learning approximately objective priors. In Proceedings of the Thirty-Third Conference on Uncertainty in Artificial Intelligence, UAI 2017, Sydney, Australia, August I 11-15, 2017. AUAI Press, 2017. (Cited on p. 10.)
Eric T. Nalisnick and Padhraic Smyth. Learning priors for invariance. In International Conference on Artificial Intelligence and Statistics, AISTATS 2018, 9-11 April 2018, Playa Blanca, Lanzarote, Canary Islands, Spain, volume 84 of Proceedings of Machine Learning Research, pages 366-375. PMLR, 2018. (Cited on p. 2.)
Rajesh P. N. Rao and Daniel L. Ruderman. Learning lie groups for invariant visual perception. In Michael J. Kearns, Sara A. Solla, and David A. Cohn, editors, Advances in Neural Information Processing Systems 11, NIPS, 1998. (Cited on p. 9.)
David W. Romero and Suhas Lohit. Learning partial equivariances from data. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Ad- vances in Neural Information Processing Systems 35: Annual Conference on Neural Infor- mation Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA. November 28 De- wwtpccppppppppeeeeeeeeeeeee fallss weiles/ppppe////eeeeeeeaaeaaeeae Saaper/202//nnn
cember 9, 2022, 2022. URL http: /twp:er.//pppsaaapaeee @c51d1fo4b007545777400 htte://ddaeeeeee66e 18w554666666611 Abst---ttttttteeee Confeeenteoooofoeccnee html. (Cited on p. 2.)
Cédric Rommel, Thomas Moreau, and Alexandre Gramfort. Deep invariant networks with differen- tiable augmentation layers. In NewrIPS, 2022. (Cited on p. 2.)
Pola Elisabeth Schwöbel, Martin Jørgensen, Sebastian W. Ober, and Mark van der Wilk. Last layer marginal likelihood for invariance learning. CoRR, abs/2106.07511, 2021. (Cited on p. 2.)
Zhixin Shu, Mihir Sahasrabudhe, Riza Alp Güler, Dimitris Samaras, Nikos Paragios, and Iasonas Kokkinos. Deforming autoencoders: Unsupervised disentangling of shape and appearance. In Vittorio Ferrari, Martial Hebert, Cristian Sminchisescu, and Yair Weiss, editors, Computer Vision ECCV 2018 15th European Conference, Munich, Germany, September 8-14, 2018, Proceedings, Part X, volume 11214 of Lecture Notes in Computer Science, pages 664 680. Springer, 2018. doi: 10. 0007/998-3-030-00 1110000149---444 40. URL https://ddo /ttp///....rr//00.000/00///----0 40. (Cited p. 10.)
664-680.
14

---

Sharvaree Vadgama, Jakub Mikolaj Tomezak, and Erik J Bekkers. Kendall shape-vae: Learning shapes generative framework. In NeurlPS 2022 Workshop on Symmetry and Geometry in Neural Representations, 2022. (Cited on p. 10.)
Tycho F. A. van der Ouderaa and Mark van der Wilk. Learning invariant weights in neural networks. CoRR, abs/2202. 12439, 2022. (Cited on pp. 2, 4, 6, and 17.)
Mark van der Wilk, Matthias Bauer, S. T. John, and James Hensman. Learning invariances using the marginal likelihood. In Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Moniréal, Canada, pages 9960-9970, 2018. (Cited on pp. 2 and 6.)
B. S. Veeling, J. Linmans, Winkens, T. Cohen, and M. Welling. Rotation equivariant cnns for digital pathology, September 2018. URL https: sttps://ddi.org//// 1077/977--3-0000000----- 24. (Cited on pp. 23 and 26.)
Mike Walmsley, Chris Lintott, Tobias Géron, Sandor Kruk, Coleman Krawezyk, Kyle W. Willett, Steven Bamford, Lee S. Kelvin, Lucy Fortson, Yarin Gal, William Keel, Karen L. Masters, Vihang Mehta, Brooke D. Simmons, Rebecca Smethurst, Lewis Smith, Elisabeth M. Baeten, and Christine Macmillan. Galaxy Zoo DECaLS: Detailed visual morphology measurements from volunteers and deep learning for 314 000 galaxies. 509(3):3966-33 3988, January 2022. (Cited on pp. and 22.)
Robin Winter, Marco Bertolini, Tuan Le, Frank Noé, and Djork-Arné Clevert. Unsupervised learning of group invariant and equivariant representations. In NeurlPS, 2022. (Cited p. 10.)
Jin Xu, Hyunjik Kim, Thomas Rainforth, and Yee Whye Teh. Group equivariant subsampling. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, vir- tual, pages 5934 5946, 2021. URL https: /tteps::/preeee http://prreeedenn..eeee neurips cc/ wttssccccpppppee httpp/pappp///////////// / 2ea6241cf767 www.cff0f00900ff 1885 Abstract. html. (Cited on 10.)
Jianke Yang, Robin Walters, Nima Dehmamy, and Rose Yu. Generative adversaria symmetry discovery. International Conference on Machine Learning, ICML, 2023, (Cited on pp. 9, 25, and 26.)
15

---

A Connections to MLL Optimization
As we will now show, Algorithm has connections to marginal log-likelihood (MLL) maximization via VAE-like amortized inference. Given the graphical model in Figure 2, we can derive an Evidence Lower BOund (ELBO) for jointly learning the generative and inference parameters with gradients:
// p p(x,, n. x) dn dx
log p (x) log
(8)
/f p(x n. x) (n x) pa (x) dn dx
log
/f pxx11 n, P) (1 8) pe (x) Iw (n, x x)
log
dn dx
(9)
qwo (n, & x)
p X 1) Pvp (1 x) pe (x)
=log
(10)
qw (1, X x)
EE logp (x n, x)] DKL 1 (n, & x) || Pyp (1 Po [X)] xw(n.x12)
(11)
diveereeeeece
likelibood
= - 1p, w);
(12)
is some generative model-e.g. e.g., VAE-fo prototypes, with parameters 0, and qw (n, x x) = Iw (n x) ( x. n). Now, we can show that the gradient of the likelihood term in (x) po(x) where Po
the ELBO is approximated by the gradient of our SSL loss on line I of Algorithm 1:
V
[log p (x | n)] lim N(x o=_0
E x, n)
(13)
lw(n | x)r((
D p 11) = Tn (x))
Ty (x),, o
V
E
[log Tn (x),
(14)
an(n x),((((,n))
take sample, n aw (n x):
V log N(x TT,(),,0²)
(15)
definition of Gaussian PDF:
Ta (.)1 /o" log (V2TT)
Va 0.5 |
(16)
D drop constant term:
V 0.5 mse (x, Tn (2)) /02.
(17)
The negative sign is due to the fact that the ELBO is maximized, whereas our SSL loss is minimized. The gradient of the KL-divergence term w.r.t. up is approximated by the gradient of our MLE loss on line 8 of Algorithm 1:
V DKL law (n. x x) || pa (n x)) pe (x))
(18)
definition of DRL:
log Sw (1 11) (2 xx, n)
E x))((1 xx
(19)
(w(n
pcp (1 x) po (x)
D drop constant terms and use & T (x):
log Pv (n T (x)]
V
E (n
(20)
x)
take sample, na qwo (1
x):
Vy log 7h Tai (z))
(21)
Note that the sampling approximations in both (15) and (21) also apply to VAE-like amortized inference algorithms.
While ELBO training and our algorithm share some similarities, some key differences exist. For in- stance, we do not learn the generative and inference models jointly. This disjoint training is equivalent
16

---

8
9
6
Figure 13: Failure of an invariant VAE encoder. Top: MNIST digits sampled from the test set. Mid: Prototypes produced by VAE who's encoder is made invariant using (22), where 1l U (-7 m and 77aa (0.25, 0.25, x, 0.25, 0.25). Bot: Reconstructed digits. The model becomes stuck in a local optima where the prototypes and 'reconstructions' are all circles and rings of various sizes depending on the input image. The averaged latent code is free of (e.g..) rotation information but has also lost almost all information that identifies each digit.
to ignoring the gradient V DKL. m9 (n, x x Pup (n {) Pa (x)] when training qw (1 x). This KI.-divergence has two components: entropy -H [u>] and cross entropy H. [gwsP%bPP Assuming that Py. (nl &) is sufficiently flexible, the cross entropy term should not have a significant impact on qw (n x) since Pp trained to match fw- On the other hand, lw (1) x) should be close to delta since there should be single prototype for each x. Thus, encouraging high variance with entropy term might actually be harmful. Another difference is that we do not need to learn pe (X), which has the benefit that we can learn the symmetries in dataset without having to learn to generate the data itself, greatly simplifying training for the complicated dataset. Furthermore, actually evaluating the gradient of the likelihood term in (12) is challenging due the fact that (xx X, n) is delta.
Given all of these differences, it might be natural to question the utility of the comparison between our algorithm and maximization of (12). Perhaps the most useful connection to draw is that of Equations (18) and (21), which motivates our MLE learning objective for Pwu (11 x) as being closely related to the process of learning a prior in an ELBO.
In an early version of this work [Allingham an ELBO similar to (12), with the main difference being that & was modeled using VAE and invariance was incorporated into the VAE encoder. We constructed an invariant encoder da (z x)
al., 2022], we trained a variant of the SGM using
from a non -ovviiaii encoder d (z x):
t (z x) = E, [90 (z x)],
(22)
following Benton et al. [2020], van der Ouderaa and van der Wilk [2022], Immer et al. [2022]. We found that this approach worked well for single transformation (e.g., rotation) but that it quickly broke down as the space of transformations was expanded (e.g., to all affine transformations; see Figure 13). We hypothesize that the averaging of many latent codes makes it difficult to learn an invariant representation z without throwing away almost all of the information in x. This further motivates our SSL algorithm for learning invariant prototypes. A similar observation was also made by Dubois et al. [2021]. who found that an SSL-based objective was superior to an ELBO-based method for learning invariant representations in the context of compression.
B Further Practical Considerations
This section elaborates on Section 3.1 and provides additional considerations.
Suitability of NN architectures. The architecture of fa must be compatible with learning an equivariant mapping from x to n. For example, a standard CNN requires many convolutional filters to represent function that is (approximately) equivariant continuous rotations [Maile et al., 2023].
X'-spae vs. H-space SSL objective. One might notice that it is possible remove the T operations from both paths of the SSL objective in Figure 4 and still have a valid objective (in H-space rather than X'space). However, the X'ssace version is preferred since different parameters 7) 722 can map to the same transformed element Tm (xx) Tnp (æ)). E.g., consider rotations transformations applied to various shapes: for a square Top = Taoe =\t Tis0 = T270 all map to the same transformed image, and an H-space objective incorrectly penalizes differences of in x 90 in 17 values.
17

---

We compare rotation inference nets nets-with hidden layers of dimensions [2048, 1024, 512, 256, 128)
trained for 2k steps using the AdamW optimizer with a constant learning rate of 3 10 and a batch size of 256-trainne on fully rotated MNIST digits using both spscaee and H-space SSL. objectives:
Objective
X-mse
T7-mse
X' d'-space
0.2387
0.9715
H-space
0.3567
0.4736
average of X'sspace and H-space
0.3129
0.4619
When using the H-space objective, we see the distance in observation (X) space.
Learning Zw (n]x) instead of fw. We found that learning fw probablistically-ie.. allowing for some uncertainty the transformation during the training process by parameterizing density over H. with geo (1 (x) and sampling --roovidees small improvements in performance. The distribution aw ()) quickly collapses to delta. Thus, we hypothesize that the added noise from sampling acts as regularizer that is helpful at the start of training.
Inference network blurring schedule. Occasionally, depending on the dataset, random seed, kind of transformations being applied, and other hyperparameteer, training the inference network fails, and the prototype transformations would be 100% lossy-iee., they would result in completely empty images regardesss of the strength of the invertibility loss. We found that we could prevent this from happening by adding a small amount of Gaussian blur to each example. Furthermore, we found that we only needed to add this blur for a small fraction of the initial training steps to prevent the model from falling into this degenerate local optima.
Averaging multiple samples for the SSL loss. Just as we found averaging the MLE loss over multiple samples improve performance, so too averaging the SSL loss.
We compare rotation inference nets-with hidden layers of dimensions [2048, 1024, 512, 256, 128] trained for 2k steps using the Adam W optimizer with cosine decayed with warmup learning rate schedule that starts at x 10-4, increases to 3 x 10-4 in 500 steps, and then decreases to x 10-7, with batch size of 256nneainne on fully rotated MNIST digits using the SSL objective averaged over 1, 3, 5, 10, and 30 samples:
Samples
X-msc
1
0,0981
3
0.0901
5
0.0840
10
0.0853
30
0.0870
As the number of samples increases, x-mse decreases until saturating around samples. Note that this relationship is not likely to be monotonically decreasing because there is random noise in each training run (i.e., due random NN initialization, etc.). That said, we expect it will decrease on average as the number of samples increases. We find 5 samples to be good trade-off between improved performance and increased compute.
Symmetric SSL loss. In our SSL loss, based on Figure 4, we are essentially comparing the prototypes given and Zmn (a randomly transformed version of æ). An alternative to compare the prototypes given Frmdl and Tmnd2 two randomly transformed versions of x::
T (amas) (FFmll) Tiu (Ze (Wmiz)) Fmall TTnan (x), Frnd2 Tryas (2:), Mmdll Mmd ~ p(1ma)-
(23)
As before, modify this loss to allow us compose transformations to get
Th (Xmee)
T--1 (Teae) (Xmma) - Xmd2
(24)
18

---

The motivation for using this 'symmetric' SSL loss is that it provides the inference network with additional data augmentation the inference network now unlikely ever to see the a twice. We find that while this works well for MNIST, it does not work well for dSprites. This is because the transformations dSprites in dSprites are more lossy than those for MNIST. E.g., it easier to shift small sprite out of the frame of an image compared to a large digit. Thus, the symmetric loss results in much higher variance when used with dSprites, which negatively impacts training.
Composing affine transformations of images. Care must be taken when composing affine trans- formations of images when implemented via a coordinate transformation (e.g., affine-grid & affine_ _aapnee_smpeeeeeeeeeee e in PyTorch, or scipy map_ coords in Jax). To compose two affine transformations parameterised by v and 772, the affine matrices T T(n)), ((n2) need to be right-multiplied with one another; in other words Typ o Tm, = T(,, (T'nn)T(nn)) ((n))- This is because, in these implementations of affine transformation of images, the affine transformation is applied to the pixel grid (i.e., the reference frame), rather than to the image itself. In effect, the resulting transformation as applied to the objects in the image is the opposite; if the reference frame moves to the right, the objects in the image move to the left, etc. More concretely, when the reference frame is affine-transformed by T, the image itself is affine- tranne-rranffomee by T
Overfitting of the generative network. While we did not observe any overfitting of the inference network (likely due the built-in 'data augmentation' of our SSL loss, and the general difficulty of learning a function with equivariance to arbitrary transformations), we did find that the generative network was prone to overfitting We addressed this by using a validation set to optimize several relevant hyper-parameters (e.g., dropout rates, number of flow layers, number of training epochs, etc.); see Appendix C.
Learning Ppp (n X) with imperfect inference, continued. To encourage p4 (n {) produce the same distribution for the inconsistent prototypes produced by Iw (n x)), we add consistency loss to line 8 of Algorithm the MLE objective:
N
Leonsistoncy (1))
log - logpsl
(25)
N2
where pı Psp (M (d)) and at is due to the 7th Und sample.
C Experimental Setup
We use jax with flax for NNs, distrax for probability distributions, and optax for optimizers. We use ciclo with clu to manage our training loops, mllcollections to specify our config- urations, and wandb to track our experiments. The code is available at https: //gg cambridge- searridggengggssgee mlg/sgm.
thub. //github.. com/
Unless otherwise specified, we use the following NN architectures and other hyperparameters for all of our experiments. We the Adam optimizer with weight decay of 1 1x10-1 global norm gradient clipping, and a linear warm-up followed by a cosine decay as a learning rate schedule. The exact learning rates and schedules for each model are discussed below. We use a batch size of 512.
All of our MLPs use gelu activations and LayerNorm. some cases, we use Dropout. The structure of each layer is Dense - gelu- + LayerNorm -> Dropout. Whenever we learn predict scale parameter o,, it is constrained to be positive using a softplus operation.
Inference network. We use a MLP with hidden layers of dimension [2048, 1024, 512, 256]. The network outputs a mean π prediction for each example and the uncertainty as mentioned in Appendix B-is is implemented as a homoscedastic scale parameter. We train for 60k steps. For each example, we average the loss over 5 random augmentations. In some settings-also mentioned in Appendix B-we add small amount of blur to the images with Gaussian filter of size 5 for the first 1% of training steps. The a value for the filter was linearly decayed from their maximum The initial maximum value is specified below.
19

---

Generative network. Our generative model is a Neural Spline Flow [Durkan et al., 2019] with We use an MLP with hidden layers of dimension [1024, 512, 512) as shared feature extractor. The base normal distribution's ' mean and scale parameters are predicted by another MLP, with hidden layers of dimension [256, 256), whose input is the shared feature representation. The parameters of the spline at each layer of the flow are predicted by MLPs with single hidden layer of dimension 256, with a dropout rate of 0.1, whose input a concatenation of the shared feature representation, and the (masked) outputs of the previous layer. For each example, we average the loss over random augmentations.
bins the range [-3,3], 3].
C.1 MNIST under affine transformations
We make use of the MNIST dataset [LeCun et al., 2010], which is available under the MIT license.
We split the MNIST training set by removing the last 10k examples and using them exclusively for validation and hyperparameter sweeps.
When randomly augmenting the inputs for our SSL (see Section 2.1 and Figure 4) and MLE (see Section 3.1) losses, we sample transformation parameters from 24( ( Mmxx). where Mmax = (0.25, 0.25, π, 0.25, 0.25) is the maximum (z-shift, y-shift, rotation, x-scale, y-scale) applied to the images. All affine transformations are applied with bi-cubic interpolation.
Inference network. The invertibility loss Livertibilty (7) is multiplied by factor of 0.1. For the VAE data-efficiency results in Figure 11, we performed the following hyperparameter grid search for each random seed and amount of training data:
blur Oinit E 30,
gradient clipping norm E [3, 10),
learning rate E 1x10
3, xxx10000x
1x10-")
initial learning rate multiplier E [3x10-2, 2x100-
final learning rate multiplier e [1x10-3.33] 113100,33x00+,, and
warm-up steps % E [0.05.0.1,0.2 0.1,
All of the other MNIST affine transformation results use a blur Gini of 0, a gradient clipping norm of 10, learning rate of 3x 10-1, an initial learning rate multiplier of x x0022, a final learning rate
multiplier of x x00-3, and warm-up steps % of 0.2, which are the best hyperparameters for 50k training examples with an arbitrarily chosen random seed. We use the 'symmetric' SLL loss discussed in Appendix B.
Generative network. We use an initial learning rate multiplier of 0.1, a gradient clipping norm of 2, and a warm-up steps % 0.2. For the VAE data-efficiency results in Figure 11, we performed the following hyperparameter grid search for each random seed and amount of training data:
learning rate E [3x10 1x3x00
final learning rate multiplier E [0.3, 0.03).
number of training steps [7.5k, 15k, 30k, 60k).
number of flow layers [4, 5, 6),
shared feature extractor dropout rate E [0.05, 0.1, 0.2], and
consistency loss multiplier& E [0, 1] (whether or not to use (25)).
Note that we use the log-likelihood of the validation data under the generative model select the best dyper-parameters Le., we do not use the total loss, which may or may not include the consistency term, since these losses are not directly comparable. We require trained inference network when sweeping over the generative network hyperparameters. We use the inference network hyperparameters for the same (random seed, number of training examples) pair. All of the other MNIST affine transformation results use a learning rate of xx final learning rate multiplier of 0.03, 60k training steps, flow layers, a dropout rate 0.2 the shared feature extractor, and consistency loss multiplier of 1, which are the best hyperparameters for 50k training examples.
20

---

C.2 MNIST under color transformations
We follow the same setup as above for color transformation on the MNIST dataset, with the following exceptions. We do not the inference and generative networks, we constrain the outputs to be in 1-7max. 7max + (0.5, O., O.), where max (0.5, 2.301, 0.51) using with tanh and scale bijectors. We randomly augment the inputs by sampling transformation parameters from 21 7-7x + (0.5, 0., 0.), Mmax + (0.5, O., 0.)).
invertibility loss when training the inference network. Instead, for both
Inference network. We use a blur mit of 3, a gradient clipping norm of 2, a learning rate of 3x10-= 4, initial learning rate multiplier of 1x10 1x12, final learning rate multiplier of 1 x 10 and warm-up steps % 0.1, which were chosen using the same grid sweep as MNIST with affine transformations.
Generative network. We use learning rate of 3 1xx00-3 with initial learning rate multiplier of 1x10-1 1000-000 a final learning rate multiplier 3x100= 2, 15k training steps, flow layers, and a dropout rate of 0.2 in the shared feature extractor.
C.3 dSprites under affine transformations
We make use of the dSprites dataset [Matthey et al., 2017], which is available under the Apache 2.0 license.
For our dSprites experiments, we follow the same setup for MNIST under affine transformations above, with the following exceptions. We do not use an invertibilityy loss when training the inference network. Instead, for both the inference and generative networks, we constrain their outputs to be in Mmaax M, where Mmax (0.75, 0.75, n. 0.75, 0.75) using with tanh and scale bijectors. We do not use the 'symmetric' SSL loss discussed in Appendix B.
Inference network. We randomly augment the inputs by sampling transformation parameters from U(- mmax Mmax), where mmax matches the constraints above. We use blur Oini of 3, a gradient clipping norm of 3, a learning rate of 1 x xx003 an initial learning rate multiplier of 3x10-2 a final learning rate multiplier of 1x10-3 and a warm-up steps % of 0.05, which were chosen using the same grid sweep as MNIST with affine transformations.
Generative network. We randomly augment the inputs by sampling transformation parameters from U1 7mux 0.75, Mmxx x 0.75), where mmux matches the constraints above. We use learning rate of 3 10-0 final learning rate multiplier of 0.3, 60k training steps, 6 flow layers, and a dropout rate of 0.05 in the shared feature extractor, which were chosen using the same grid sweep as MNIST with affine transformations.
Although we swept over the consistency loss multiplier, we accidentally always used a consistency loss multiplier of in our experiments. This means that for some (random seed, amount of training data) pairs the performance of our generative network is slightly lower than should be since the chosen hyperparameters may correspond to consistency loss multiplier of O. We include this detail for reproducibility but note that does not change our findings any material way.
C.3.1 dSprites Setup
The original dSprites dataset contains sprites with the following factors of variation [Matthey et al., 2017].
Color: white
Shape: square, ellipse, heart
Scale: values linearly spaced in [0.5,
Orientation: 40 values linearly spaced in [0, 2m]]
X position: 32 values linearly spaced in [0, 1]
Y position: 32 values linearly spaced in [0, 1]
21

---

position
position
6000
200
000
400
1.000
200
200
10000
600
10000
ellipse
400
5.00
200
6.00
10000
4000
750
730
heart
2000
250
Figure 14: Latent factor distributions for our modified dSprites data loader.
The dataset consists of sprites with the outer product of these factors, for total of 737280 examples. We modified our data loader to resample the sprites proportional to the following distributions on the latent factors conditioned the shape.
square
Scale: TruncNorm (N 0.75, o² -0.2, min 0.55, max- 1.0)
Orientation: 4(0.0, 2T)
position: 2/(0.5, 0.95)
position: (0055, 0.95)
ellipse
Scale: TruncNorm (0.65, 0.15, 0.5, 0.85)
Orientation: (0.00, =/2)
X position: TruncNorm (0.5, 0.25, 0.1, 0.9)
-Y position: TruneNorm (0.5, 0.15, 0.35, 0.65)
heart
Scale: (009,, 1.0)
Orientation: 5 (0.0)
position: 20011, 0.5)
Y position: 0.5 00.1, 0.3) 0.5 ((.77, 0.9)
An example of the resulting empirical distributions over the latent factors is shown in Figure 14, The three shapes are sampled with equal proportions.
C.4 GalaXyMNIST under affine and color transformations
We make use of the GalaxyMNIST dataset [Walmslcy et al., 2022], which is available under the GPL-3.0 licence.
For our GalaxyMNIST experiments, we follow the same setup as for MNIST under affine transformations above. with the following exceptions. We do not use an invertibility loss when training the inference network. Instead, for both
inference and generative networks, we constrain their outputs to be in Mmax Mmxx + (0., 0., O., O., 0., 0.5, O., O.), where Max (0.75, 0.75, n, 0.75, 0.75, 0.5, 2.31, 0.51) using with tanh and scale bijectors. This dataset con- tains 10k examples. We use the last 2k test set, and the previous 1k validation set.
22

---

Inference network. We use a MLP with hidden layers of dimension [1024, 1024, 512, 256]. We train for 10k steps. We randomly augment the inputs by sampling transformation parame- ters from (m77xxx + (0., O., O., O. O., 0.5, 0., 0.), mmxx + (0., 0., 0., 0., 0.5, O., 0.)), where mmmx matches the constraints above. For the VAE data-efficiency results in Figure 12, we performed the same hyperparameter grid search as above for each random seed and amount of training data. All of the other Galaxy MNIST results use a blur iint of 0, a gradient clipping norm of 10, learning rate MalaxyMNIST of 3x 10-4, an initial learning rate multiplier of 1x 10 and a warm-up steps % 0.2, which are the best hyperparameters for 7k training examples with an arbitrarily chosen random seed. We use the 'symmetric' SLL loss discussed in Appendix B.
final learning rate multiplier of 3 10-1,
Generative network. We randomly augment the inputs by sampling transformation parameters from U(-7max 0.75 + (0., O., O., 0., 0., 0.5, O., O.), Mmux x 0.75 (0., 0., 0., 0., 0., 0.5, 0., 0.)), where 7max matches the constraints above. For the VAE data-efficiency results Figure 12, we perform the same hyperparameter grid search above for each random seed and amount of training data, with the following changes." The sweep for number of training steps is [3.75k, 7.5k, 15k]. All of the other GalaxyMNIST results use learning rate of 3 10 4, a final learning rate multiplier of 0.03, 15k training steps, 4 flow layers, a dropout rate of 0.05 in the shared feature extractor, and a consistency loss multiplier of 1, which were chosen using the same grid sweep for an arbitrary random seed and 7k training examples.
C.5 PamelCon under affine and color transformations
We make use of the PatchCamelyon dataset [Veeling et al., 2018), which is available under the Creative Commons Zero v1.0 Universal license.
We resized the images from 96 x 96 pixels to 64 x 64 using bilinear interpolation. The dataset has dedicated train, test, and validation splits which use without any modifications.
We follow the same setup as for GalaxyMNIST under affine and color transformations above, with the exceptions listed below. We only used single random seed.
Inference network. We train for 20k steps.
Generative network.
The sweep for number of training steps is [15k, 30k, 60k).
C.6 VAE, Aug and VAE
Our VAEs use latent code size of 20. The prior is a normal distribution with learnable mean and scale, initialized to Os and Is. respectively.
Our VAE encoders are LeNet-style CNNs with convolutional feature extractors followed by MLP with single hidden layer of size 256. The convolutional feature extractors use gelu activations and LayerNorm. The structure is Conv gelu LayerNorm. All Conv layers use 3x3 filters. The first two Conv have stride of 2, while all others have a stride of 1. In between the convolutional layers and the MLP, there special dimensionality reduction Conv with only filters followed by f flatten. For each dimension of the latent code, the encoder predicts a mean μi and a scale a. The means and scales initialized to Os and 1s, respectively.
Our VAE decoders are inverted versions of our encoders. That is, we reverse the order of all of the Dense and Conv layers. The dimensionality reduction Conv layer and the flatten operation replaced with the appropriate Dense layer and reshape operation. We replace all other Conv layers with ConvTransposed layers For each pixel of an image, the decoder predicts a mean μ. We learn homoscedastic per-pixel scale a. The scales are initialized to 1.
"Our GalaxyMNIST results have the same issue as our dSprites results the sweep included consistency loss multiplier which was always set to value of in our experiments. This results in some small performance degradations.
*Our PatchCamelyon results have the same consistency multiplier issue our dSprites and GalaxyMNIST results.
23

---

We use an initial learning rate multiplier of 3x 3x10-2 and final learning rate multiplier of 1x10 We run the following grid sweep for each (random seed, number of training examples, maximum added rotation angle) triplet:
learning rate E [3x10-3, 6x 66x10-3 9x 10 10-
convolutional filters [644, 128), (64, 128, 256)],
number of training steps E [5k, 10k, 20k), and
warm-up steps % E [0.15, 0.2].
When running the sweep for Vug and Inv VAE, we use the inference and generative network hyperparameters for the same (random seed, number of training examples) pair.
C.6.1 PatchCamelyon
For our PatchCamelyon experiments, we use only a single random seed and a slightly modified hyperparameter sweep:
learning rate E [3x10-3,6x 16-10-
convolutional filters [644, 128), (64, 128, 256), (128, 256, 512)],
number of dense hidden layers [1, 2].
number of training steps E [20k, 30k, 40k], and
warm-up steps % E [0.15].
C.7 Parametrisations of Symmetry transformations
We consider five affine transformations: shift in x, shift in y, rotation, scaling in x, and scaling in
We represent these transformations using affine transformation matrices A = exp (2: 1), G.), where G; are generator matrices for rotation, translation, and scaling; see Benton et al. [2020]. The transformations are applied to an image by transforming the coordinates (r, y) of each pixel, as in Jaderberg et al. [2015]: [x' y/ 1]" A [Ir y/
To parameterize color transformations, we use an equivalent representation color images in Hue- Saturation- Value (HSV) space, where each pixel is represented
tuple (h, v) {1-x, π| x [0, 1] x [o, 1]}. Intuitively, HSV space represents the color of each pixel in a conical space where the hue corresponds to the rotation angle around the cone's 's vertical axis, the saturation corresponds to the radial distance from the cone's 's center, and the value corresponds to the distance along the cone's vertical axis, with value of 0 corresponding the tip of the cone, and value of corresponding to the base of the cone. We color-transform an image by transforming each pixel as
(h 2mT)) mod 2Tt max (0, min(s cxp(n)))
(26)
max(0, min(c cxp()))),
We therefore obtain r) = (7h, Me 7o. E {0, 1] x RR R). We choose this specific form of parametriz- ing the 77 parameters order to gain the convenience of simply adding and subtracting in n space when carrying out color transform compositions and inverses. More concretely, with our chosen parametrization, we obtain the property that Tr, Typ +7. Therefore, we can easily perform compositions and inversions in v space for color transformations without resorting to matrix multipli- cations. In order to achieve this, we first consider hue, which easy to parametrize in an additive fashion using modulo operation due to the fact that hue is represented space. On the other hand, saturation and value are discontinuous parameters that vary between and and cannot be directly modeled in additive fashion, as they can't take values outside their range. Instead, we model them as multiplicative factors + where we first exponentiate T] and nh to ensure the multiplicative factors positive. We further clip the obtained values to ensure they are in the range [0, 1]. This parametrization allows us to effectively add parameters to compose them, as the multiplicative factors compose in exponent space.
rotation angle in HSV
In order to ensure that we can easily backpropagate through the clipping operation, we define passthrough_ ough_clip function in Jax, where we define custom gradient that doesn't zero out gradi- ents even if the inputs the function are out bounds. We find that using the operation is essential training the model.
24

---

D Compute Requirements
The experiments for this paper were performed cluster equipped with NVIDIA A GPUs. All model training requires only a single such GPU. However, we used up to 64 GPUs at a time to nyper-parameter searches in parallel. Including exploratory experiments, all hyperparameter sweeps, discarded runs, etc., the total compute used for this paper is approximately 250 A100 GPU days. The total cost reproduce the experiments in the paper is approximately 135 A100 GPU days. We break this cost down as follows. Note that the cost for different figures do not naively sum as hyper-paramett sweeps for some figures are reused for others, as discussed Appendix C.
Figure Sa: 6 days
Inference net sweeps: 4 days
Generative net sweeps: 2 days
Figure 8b: 3 days
Inference net sweeps: 2 days Generative net sweeps: day
Figure 8c: 3 days
Inference net sweeps: 2 days Generative net sweeps: day
Figure 8d: 7 days
Inference net sweeps: 6 days Generative net sweeps: day
Figure days
Inference net sweeps: 2 days Generative net sweeps: I day
Figure 10: 2 days
Inference net sweeps: 2 days
Figure 11: 69 days
Inference net sweeps: 30 days Generative net sweeps: 12 days
VAE sweeps: 27 days
Figure 12: 53 days
Inference net sweeps: 36 days Generative net sweeps: 8 days
VAE sweeps: 9 days
E
Additional Results
E.1
Comparisons to LieGAN
In this section, compare the ability of our method to learn symmetries LieGAN [Yang et al., 2023], which uses cenerator-discrimnnat from dataset using generative adversarial training. Similar to [Yang al., 2023], transform the MNIST dataset to have rotations in the range [-450, 45%), which ensures the dataset contains SE(2) symmetry (rotations and translations). The dataset is processed and our method is trained as described in Section 4.1. For LieGAN, following the experimental design of (Yang et al., 2023], we set the number of generator channels the generator model. The discriminator model consists of pre-trained LeNet5 feature extractor as the backbone, and the validator is 3-layer MLP with 512 hidden units and ReLU activations, We train the GAN for 100 epochs with a batch size of 64, and obtain the Lie matrix below
framework to automatically discover equivariances
and consider learnable 6-dimensional Lie matrices in
0.02
-0.4
0.28
L
0.33
0.08
-0.05
0
25

---

1
5
6
6
5
Figure 15: Learnt augmentation distribution for the MNIST dataset rotated in the range
[-45%, 45%] for our SGM model, and the LieGAN method. The columns correspond to distri- butions for translation 2, translation in 3/, rotation, scaling in z, and scaling in y/. (Row 1-5) Our SGM learns accurate ranges of rotational invariance present in the training dataset of a width of 2/2 for most training examples, along with learning the natural invariances present in the training data for translations and scaling. Furthermore, for certain digits (i.e. O), the SGM model accurately predicts uniform distribution from -, T. signifying that rotationally invariant digits such as O would not display
narrow rotational invariance. (Row 6) On the other hand, the LieGAN model learns a single Lie matrix across the entire training dataset that encodes the maximum possible range of transformations, and predicts uniform distribution between those ranges. that LieGAN inaccurately predicts large range for translations in a, and does not recover the correct range of rotational invariances present in the training dataset.
be
In Figure 15, we can that LieGAN struggles correctly recover the range of invariances present in the training dataset, especially for translations in æ. It is also unable to provide a fine-grained representation invariances depending on specific examples or type of digits. We note that we re-implemented the rotated MNIST experiment from Yang et al. [2023]. as the code for the image domain experiments was not open-source. Hence, the choice of using pre-trained LeNet5 model for the discriminator, and the specific hyperparameter configurations, were informed decisions made by us based on ablations. However, our results appear to inline with those presented by Yang et al. [2023]; concretely, we note that the results presented in their paper also display a mismatch between the invariances present the dataset and those learned by LieGAN. For example, in their Figure 11, we see that the sampled digits often rotated by significantly more than 45%. Furthermore, evidence of typical GAN mode collapse, with many very similar rotations for each digit.
E.2 PatchCamelyon - Boundary Effects
In this section, we provide "negative" result for our SGM when applied the PatchCamelyon dataset [Veeling et al., 2018]. The examples in this dataset, unlike those used in Section 4, contain "content" up the boundaries of the images.
Figure 16 shows examples of the prototypes and learned distri- butions for this dataset, with affine and color transformations. In particular, the allowed rotation was between + 180%, while the ac- tual dataset has only rotational invariance of +n 90", We see that in some cases the prototypes are rotated by close to +n x 45"
-1000
IWLB
relative to the original images. In other cases, the rotation of the prototypes relative to the original images is closer to tn x 90".
In the latter case, the learned distribution over rotation is close to the true distribution, but in the former case, the model learns a distribution that closer to uniform. As a result, the resampled digits often display boundary effects that are not present in the original dataset. Otherwise, our SGM has learned reasonable distributions for translation, scaling, and HSV transformations.
-3000
16384
65536
262144
Num. Train
Figure 17: VAE data-efficiency for PatchCamelyon.
26

---

(a) Top: samples from the Mid: prototypes for each example. Bot: resampled versions of each example given the prototype.
0
(b) From left right, test examples, their prototypes. and the corresponding marginal distributions translation x, translation y, rotation, scaling in ir scaling in y, hue, saturation, and value.
Figure 16: Prototypes and learned distributions for PatchCamelyon.
0.0300
Num. train
0.0275
12500
37500
Recon, MSE
0.0250
1200
TWLB
0.0225
0.0200
Model
0.0175
IAE
AugVAE
WAE aug
0.0150
800
90
180
90
180
Maximum Added Rotation Angle (*)
Figure 18: Incorporating symmetries improves data efficiency. Importance-weightel lower bound (IWLB) and reconstruction MSE (mean and std. err. over 3 random seeds) for rotated MNIST with standard VAE (with and without standard data augmentations) and two VAE variants that incorporate symmetries via our SGM. Improved data efficiency is demonstrated by better performance with less training data, and reduced sensitivity to added rotations.
Figure 17 compares a standard VAE with AugVAE, an SGM-VAE hybrid model. We see that for small amounts of data, the VAE and Aug VAE perform similarly. However, as the amount of training data increases, the VAE performs better. This is likely because the SGM has not learned the true distribution over rotations.
This "negative" result highlights the importance of correctly choosing the prior transformation distributions in some settings. In this case, the performance of the SGM would have been improved by choosing a categorical distribution over rotations.
E.3 Additional Experiments
In this section, we provide additional plots to supplement those in Section 4.
Figure 18 extends the results in Figure 1l to by including additional metric: reconstruction MSE. Our findings with IWLB are consistent for this metric.
Figure 19 expands on Figure 8b in two ways. Firstly, it makes it clear that our inference network able to provide the same or very similar prototype for observations in same orbit. Secondly, it provides many more resampled examples of each digit, further demonstrating that our SGM has correctly captured the symmetries present the dataset. Figure expands on Figure 8c in the same way.
27

---

Figure 21 extends Figure 9 by including all of the digits shown in Figure 19. The conclusions are much the same before. We see that the learned distributions all make sense, especially for the most easily interpretable transformation parameter, rotations. Again, we note that smaller and bigger prototypes have appropriately different scaling distributions. Figure 22 provides the learnt marginal distributions for the digits in Figure 20. Here, we manually controlled the distributions over hue and saturation when loading the dataset, so we know that the range of the hue distribution should be approximately while the range of the saturation distribution should be around 0.3. We see that this indeed the case. However, given that most (non- moon-black) pixels are bright (i.e., close to makes sense that SGM leams multiplicative values closer to
not control the value the images, so it is more difficult to interpret those.
Finally, Figure 23 extends our dSprites results in two ways. Firstly, it provides many more resampled sprites, which also serves to demonstrate further that our SGM has captured the symmetries correctly. Secondly, the figure includes empirical distributions of positions of each of the classes of digits, which we have carefully controlled as described in Appendix C.3.1. These empirical distributions for the dataset are compared with empirical distributions for our resampled sprites. We see that although the resampled densities don't match the original densities perfectly, their general shapes and ranges are correct.
28

---

I
2
2
2
Y
C
2
2
2
2
2
2
2
C
2
Y
A
7
7
7 2
7
2
7
7
7
7
4
7
2
2
2
2
2
7
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
3
2
3
3
3
3
3
33
3
3 3 3 3 3
7
3
3
3
5
3
3
3
3
3
3
&
Σ Σ
3 3 3 & 3 3 3
3
Σ Σ
3
3
3
3
3
3 3 3
3 3 3 3 3 3 3 3 3 3 3 3 3 3
h
h
4
4 4 4 4 4 h
X
4
>
< h
4
4
6
4
4
A
4
4
4 4 4 4 4 4
A
4 6
4 4
4
A
4
4
4
4
4
4
7
4 4 4
4
4
4
9
5
5
5
S
9
5
5
5
5
5
5
5
9
5
5
6
S
9
5
5
S
S
S
S
S
S
9
S
9
9
S
5
S
S
9
S
5
5
5
5
9
5
5
5
5
5
9
5
5
5
5
5
9
5
5
5 6
5
5
9
6
6
6
9
6
6
6
6
6
6
6
9
6
6
9
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
9
6
6
6
6
6
9
6
6
6
6
6
9
6
6
6
6
6
6
4
7
3
7
^
3
4
2
7
7
7
7
7
7
V
3
7
7
4
2
2
4
4
4
4
4
4
4
4
4
4
4
4
2
4
2
7
7
7
0
7
7
3
7
2
7
7
7
>
2
2
7
7
>
2
7
8
4
8
6
S
B
8
8
4
4
8
3
44
9
0
B
3
8
8
R
8
B
8
8
3
8
8
3
8
8
8
8
8
8
8
8
8
8
8
S
8
S
8
8
8
S
S
S
8
8
8
3
8
8
6
5
9
9
0
6
9
7
7
9
7
7
6
6
9
9
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
9
9
9
6
9
9
9
9
9
6
9
9
7
7
9
6
9
9
9
9
9
Figure 19: Columns from left to right: only rotation, only translations, translation rotation scaling. Each of the blocks in this figure follows the same format. Top: examples from the same orbit. Mid: The corresponding prototypes. Bot: Resampled versions of the digits, given the prototypes.
29

---

/
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2-2
2
2
2
2
1
2
2
2
2
2
2
12222
2
2
2
2
2
122211
22222
2
2
2
2
2
2
2
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
4
4
/ / / / / / / / / / /
4
4 4 4
4
4
4
4
4
4 4 4 / / //
4/
/
/
/ / 4 4
4
4
4
4
4
4
4
4
4
4/ 4 4 4 4
4
4/ 4
4 4 4
4
4
4 4 4 4 4
5
5
5
5
5
5
5
5
55545
5
5
5
5
5 5
5
5
5
5
5
5
5
5
5
5
5
55
5
5
5
555555
555559
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
7
7
7
7
7
7
7
7
7
7
7
8
8
8
8
8
8
8
8
2
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
3
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
8
9
9
9
7
9
9
8
9
7
7
9
9
7
7
7
7
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
9
7
9
9
9
9
9
9
9
9
9
9
9
9
Figure 20: Columns from left to right: only hue, only saturation, only value. Each of the blocks in this figure follows the same format. Top: 7 examples from the same orbit. Mid: The corresponding prototypes. Bot: Resampled versions of the digits, given the prototypes.
30

---

/
2
7
2
2
3
3
3
4
4
4
6
5
5
5
6
6
6
7
7
7
Z
2
8
8
8
9
6
9
Figure 21: From left to right, test examples from MNIST, their prototypes, and the corresponding
marginal distributions over translation 2, translation y, rotation, scaling in æ, and scaling in y.
31

---

Sprototype exp(1)
Uprototype exp(1)
0
0
/
/
2
2
2
2
3
3
3
3
4
4
4
4
5
5
5
5
6
6
7
7
7
7
7
7
8
8
8
8
9 9
9
9
0.5
Figure 22: From left to right, test examples from MNIST with added hue in the range 0.67, and saturation scaled by factor in 0.6 to 0.9, their prototypes, and the corresponding marginal distributions over hue, saturation, and value.
32

---

Figure 23: From left to right, samples from dSprites, the empirical distribution over the positions of the sprites, sprites resampled using our SGM, and the empirical distributions over the resampled sprites' positions. We see that the resampled sprites are visually very similar to the original sprites in terms of sizes, rotations, and positions. Furthermore, terms of ranges, although they are imperfect in density.
that the empirical distributions match in
33

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: In this paper we present novel generative model of symmetry transformations. In our abstract and conclusion make two claims about this model: (1) it can accurately capture the symmetries in a dataset, and (2) when combined with a standard generative model we see improvements in data-efficiency. We believe that both of these claims reflect the paper's 's contributions well. In the introduction, we also discuss some aspirational goals for disentanglement and scientific discovery, however, we are clear that these are not the focus of the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Throughout the paper we provide footnotes to clarify the scope of our claims and point out their limitations (e.g.. footnote clarifies that our generative model does not always match the true generative process of the data). We also provide detailed list of potential issues when using our method in practice. Furthermore, in our conclusion, we note that our method only learns upproximate symmetries and requires super-set of possible symmetries in the data to be specified. Finally, we provide some "negative results" in Appendix E.2, which are also mentioned as a limitation in our conclusion.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-texx system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
34

---

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification: This paper contains no theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide a clear algorithm description (Algorithm 1), discussions of all of the practical issues encountered when implementing our method (Section 3.1 and Appendix B, and detailed experimental setup descriptions srrittiions-ncludiing dataset splits, model architectures, hyper-parameter settings and sweeps, transformation parameterisations, and list of software libraries used (Appendix C).
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
35

---

(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g.. a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We have provided link GitHub repository. We have not given detailed instructions for reproducing the experiments, however, all of our configurations and training seripts are provided.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co cc/
wwwpaabmisssioooool1soooo iolloccy) for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: / //nips. wwtppcccccgguuiecueeccccc ceplliccguiier gttp///ieee//eee//ee seedidde///ddbeiiii "CooeSSumisssiii siiniiissiiiiiciiii for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
We provide detailed experimental setup descriptions--incuudin dataset splits, model architectures. hyper- parameter settings and sweeps, transformation parameterisations, and a list of software libraries used in Appendix C.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, or as supplemental material.
36

---

7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: For all quantitative results, we report the mean and standard error over 3 random seeds.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative етгог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: See Appendix D for estimates of the compute costs, in the form of A100 GPU days, for the whole project as well as each of the figures in the main text.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: http///neuipppp cc httpp//bcct.tthchccccccceaaaaaee
Answer: [Yes]
Justification: We have read and acknowledged the NeurlPS Code of Ethics. We believe that our paper conforms with this code in every respect.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
37

---

If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work is foundational research that is not tied to any particular application for which we see direct path negative applications.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative socictal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymeett of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification: Our work does not pose such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
38

---

Answer: [Yes]
Justification:
We eite and provide licenses for all of the datasets used in this paper.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi paperswithcode thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors encouraged to reach out the asset's ' creators.
13. New Assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: Our paper does not release any assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and sereenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
We did not make use any crowdsourcing or human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
39

---

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We did not make use of any crowdsourcing human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
40


Reviews:
[{"review_id": "BbKI6NtnLi", "reviewer": "Reviewer_rf86", "summary": "The paper proposes a generative model of symmetry transformations. The work leverages recent parameterizations based on group theory to define a generative model, in which relaxed symmetry becomes a latent variable over which inference can be performed.", "questions": "1. The connection in App. A. offers a nice connection to the objectives used in some of the prior work. It references earlier approaches by the authors [Authors, 20XX] in which the ELBO is directly optimized and hypothesizes that ‘ averaging of many latent codes makes it difficult to learn an invariant representation a without throwing away all the information in x’ . Could the authors elaborate a bit more on this hypothesis / is this backed by any experiment?\n2. Number of samples. The x-mse in number of samples is optimal for 5 samples. Don’t we expect this table to be monotonically decreasing in number of samples?\n3. Overfitting on p(\\eta | x). For very flexible distributions of transformations, isn’t there a risk of overfitting on the parameterization? Please correct me if I have missed something in the method which counteracts this.", "limitations": "The paper offers a strong contribution proposing a probabilistic generative model that describes that as coming from transformed latent prototypes.", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "The paper is very well written and proposes an elegant method that leverages group theoretical framework to mathematically define approximate equivariances, combined with a probabilistic approach to guide discovery of symmetry transformations and learn a prototype.\n\nLastly, the reviewer notes that the paper is very well-written and has beautiful illustrations which guide the intuitive understanding of the generative model and the concept around learning a prototype.\n\nPracticality\nIt seems that the proposed SSL objective has computational benefits (easy to scale) as well as benefits in performance (ELBO objective only worked for rotations). Is this true, or am I reading this too positively? What would be potential disadvantage of choosing such loss?\n\nExperiments\nThe paper provides very thorough experimental validation on a wide range of datasets. The appendices provide extensive details on the experimental set-ups making them highly reproducible.", "weaknesses": "Intuition behind the objective.\nExperimentally the paper demonstrates that the proposed SSL objective is very effective. Apart from the computational benefits of such approach, it also seems to improve overall performance (optimizing ELBO only worked for rotations). To me it is not entirely clear why this would be the case, and it would be interesting to provide some more explanations on this - if known, of course. The first appendix was very helpful in providing context in relation to directly optimizing the ELBO.\n\nTypes of transformations.\nThe paper does not provide a lot of discussion on how the density of transformations is parameterized. In case of rotations, how is the normalizing flow constrained to remain smooth in the Lie algebra? Is the approach mostly targeted to simple (e.g. affine) groups?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> Practicality It seems that the proposed SSL objective has computational benefits... Is this true, or am I reading this too positively? What would be potential disadvantage of choosing such loss?\n\nYour understanding is correct. We provide additional discussion below. Regarding disadvantages, the SSL objective does have a few pathologies and potential ‘gotchas’ that are not present in the ELBO setting. E.g., see “Partial invertibility” in Section 3.1 and Appendix B. \n\n> The connection in App. A. offers a nice connection to the objectives used in some of the prior work. It references earlier approaches by the authors [Authors, 20XX] in which the ELBO is directly optimized and hypothesizes that ‘ averaging of many latent codes makes it difficult to learn an invariant representation a without throwing away all the information in x’ . Could the authors elaborate a bit more on this hypothesis / is this backed by any experiment?\n\nIn trying to scale that method, we observed that as we added additional transformations or increased the range of those transformations (e.g., increasing rotation from 0-pi to 0-2pi) that performance (as measured by the ELBO and reconstruction loss) degraded to the point that when using the 5 transformations applied to MNIST in this paper, the model was unable to reconstruct the digits at all. Instead it became stuck in a local optima in which the ‘reconstructions’ were all circles and rings of various sizes depending on the input image. (We’ll include some of these figures in the camera-ready version.)\n\nIn other words, the averaged latent code was successfully throwing away (e.g.,) rotation information but was also throwing away all of the information that actually identified each digit. This led us to the hypothesis that averaging of latent codes makes it difficult to learn representations that only throw away the symmetry data.\n\nOur current method is directly motivated by this observation – we aimed to develop an algorithm that could produce invariant representations without latent-space averaging. The success of our SSL objective is indirect evidence to support our hypothesis. \nThat said, our SSL algorithm has another advantage over ELBO learning – it decouples learning an invariant representation of x from reconstruction of x. That is, for ELBO learning to successfully learn an inference network, one ultimately needs a good generative network (and vice-versa). However, learning a network to generate examples given latent codes is a challenging inverse learning task. This is an observation that was also made by Dubois et al. [2021], who found that an SSL based objective was superior to an ELBO based method for learning invariant representations in the context of compression.\n\n> Number of samples. The x-mse in number of samples is optimal for 5 samples. Don’t we expect this table to be monotonically decreasing in number of samples?\n\nThanks for the question, we will clarify this in the camera-ready version. The table is not likely to be *monotonically* decreasing, due to the fact that there is random noise in each training run (i.e., due random NN initialization, etc.). That said, we would expect that it will decrease on average as the number of samples is increased. We chose 5 samples not because it provides the lowest loss, but rather because it provided a good trade-off between lower loss and increased compute cost.\n\n> Overfitting on $p(\\eta | x)$. For very flexible distributions of transformations, isn’t there a risk of overfitting on the parameterization? Please correct me if I have missed something in the method which counteracts this.\n\nIn practice, for the inference network $p(\\eta | x)$ we found that there were no issues with overfitting (the more expressive the network and the longer we trained, the better we found the test-set performance became). This is likely due to two things: (1) our SSL loss, which has ‘baked-in data-augmentation’ in the form of random transformations applied to x, and (2) learning a function with equivariance to arbitrary transformations is hard.\n\nOn the other hand, for the generative network $p(\\eta | \\hat{x})$, we did observe overfitting. We addressed this by using a validation set to optimize several relevant hyper-parameters (e.g., dropout rates, number of flow layers, number of training epochs, etc.). \n\n\n> Rating: 8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.\n\nPlease let us know if you have any remaining concerns.", "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": "", "extra_scores": {}}, {"review_id": "zmP1NGDWR1", "reviewer": "Reviewer_TzHS", "summary": "The paper proposes a generative model that disentangles the latent space into a group-invariant part -- the latent for the prototype -- and another part which represent a group element that can be applied to the prototype to reconstruct the input. A key novelty is to simultaneously learn to predict a distribution for each input over the group elements based on data. This input-dependent distribution, in principle, can be then used to generate new data that better aligns with the true group distribution.", "questions": "Please address the weaknesses I have listed above.", "limitations": "I don't think the paper mentions the limitations explicitly, which the authors should try to.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The overall architecture seems interesting and sound and the idea of learning probability distribution over the group elements is especially interesting. The results show that the distributions are dependent on the input and are meaningful for the datasets used in the experiments.", "weaknesses": "Some parts of the paper are confusing to me. \n\nWhat is the architecture for the part that predicts the distribution over the group elements? I saw one or two mentions of normalizing flows, but that is not enough to understand the details. The figure is quite unclear about it. Shouldn't the input image be also an input to the network that predicts $p_\\psi(\\eta|x)$? Also, the paper does not seem to have any information on the loss function used to train the network that predicts $p_\\psi(\\eta|x)$. A lot more clarity is needed for understanding these important details. \n\nI feel that the exact new contributions of the paper are a little unclear given many previous works doing similar things. Many of them are also mentioned by the authors in the appendix, but the contrasts with this paper are not clear. My understanding is that it is a generative model and can find the right distribution over the group elements directly from data, but there are earlier works also looking into these aspects. Some discussion here would be useful.\n\nXu et al., Group Equivariant Subsampling\nRomero and Lohit, Learning Equivariances and Partial Equivariances from Data\nShu et al., ,Deforming autoencoders: Unsupervised disentangling of shape and appearance\nChau et al., Disentangling images with Lie group transformations and sparse coding\n\nEven earlier work by Grenander, Mumford and others on Pattern Theory discusses generative models and probability distributions over groups, but these are not based on neural networks.  \n\nI think there should be at least one dataset which is a little more challenging like an image recognition dataset of more natural images. This can also help in understanding the limitations of the method. Another experiment showing the usefulness of the generative model being able to generate samples that respect the data distribution is also going to be useful.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> What is the architecture for the part that predicts the distribution over the group elements? \n\nIn short, we use an MLP with hidden layers of dimension [1024, 512, 512] as a shared feature extractor. These shared features are fed into another MLP hidden layers of dimension [256, 256] that outputs a mean and a standard deviation. The shared features are also used by MLPs with a single hidden dimension of size 256 that output the flow parameters at each layer of the neural spline flow, which has 6 bins in the range [-3, 3].\n\nPlease see Appendix D for further details about the exact NN architectures used for both the inference and generative networks. If you would like, we can include some of the details in the main text (given the extra page allowed for the camera-ready version). If so, please let us know what you would find most useful.\n\n> Shouldn't the input image be also an input to the network that predicts\n\nWe are not sure which network you are referring to. The inference network does indeed take the original image as an input. On the other hand, the generative network, which represents the distribution of $p(\\eta|\\hat{x})$, must take $\\hat{x}$, rather than $x$ as input. One way to think of why it must take $\\hat{x}$, rather than $x$, is that we want to learn the distribution of transformations present in the whole dataset. This distribution should not change depending on the transformation applied to an individual input (e.g., the rotation of a specific digit) but rather an aggregate of transformations applied to all the digits of the same type. Since the distribution is capturing a “dataset-level” transformation, it should depend on a representation of data that is invariant to those transformations (i.e., the prototype $\\hat{x}$).\n\n> Also, the paper does not seem to have any information on the loss function used to train the network that predicts\n\nWe are not sure which network you are referring to. In section 2.1, under “Transformation inference function” we discuss how the inference network  is trained. In short, we use a SSL loss depicted in Figure 4 (and equation 6). Similarly, under “Generative model of transformations” we discuss how the generative network is trained. In short, we use the inference network to generate data and we then fit the conditional flow with maximum likelihood. Both loss functions and training methods are summarized in Algorithm 1.\n\n> A lot more clarity is needed for understanding these important details.\n\nPlease let us know if there are any other clarifications required.\n\n> I feel that the exact new contributions of the paper are a little unclear... \n\nThank you for this feedback. We will update the related work sections to clarify this better. We provide a short summary below.\nWhile the goals of our paper–namely, (1) unsupervised learning a distribution over arbitrary symmetry transformations present in a dataset, and (2) leveraging this distribution for improved data-efficiency in deep generative models–and the techniques we employ to do so–(A) our SSL objective for learning invariant representations and (B) maximum likelihood learning of flexible flow models for the distributions over *partial*-symmetries–are similar/related to a wide range of existing work, to the best of our knowledge ours is the only work that incorporates all of 1, 2, A, and B. \n\nRelated methods tend to not learn the distribution of interest (e.g., [Winter et al., 2022] and “Deforming autoencoders”), do not learn invariant prototypes (e.g., [Yang et al., 2023]), focus on the discriminative rather than generative setting learning setting (see below), or construct deep-generative models for specific symmetries (e.g., [Kuzina et al., 2022] and [Vadgama et al., 2022]).\n\n> Xu et al., Group Equivariant Subsampling Romero and Lohit, Learning Equivariances and Partial Equivariances from Data Shu et al., ,Deforming autoencoders: Unsupervised disentangling of shape and appearance Chau et al., Disentangling images with Lie group transformations and sparse coding\n\nThank you for providing these additional items of related work. We will happily include them in our camera-ready manuscript. We provide brief discussions for each paper in the global comment above.\n\n> Even earlier work by Grenander, Mumford and others on Pattern Theory discusses generative models and probability distributions over groups, but these are not based on neural networks.\n\nIf you provide specific missing references, we’d be more than happy to include them in our discussion\n\n> ... at least one dataset which is a little more challenging like an image recognition dataset of more natural images... experiment showing the generative model being able to generate samples that respect the data distribution...\n\nPlease see our general response, in which we provide detailed motivation as to why our experiments are sufficiently interesting and informative. In short, we have already gone beyond many published related works, our choices of datasets cover a range of dataset sizes and data dimensionalities, and the GalaxyMNIST dataset contains high-dimensional natural images while providing challenges due to its small number of observations and large number of symmetries.\n\nPlease see Appendix F.2 for additional experiments showing that the generated samples respect the data distribution.\n\n> I don't think the paper mentions the limitations explicitly...\n\nWe have provided several limitations, e.g., in footnote 1 we discuss how our generative model might not always match the true generative process, and in our conclusion we discuss our need to pre-specify a super-set of possible symmetries.\n\nIn the camera ready version, we will also include a discussion of potential limitations of our method due to boundary effects (see our discussion with reviewer d33o). \n\n> Rating: 5: Borderline accept\n\nPlease let us know if we have successfully addressed your concerns. If so, we would appreciate it if you would consider increasing your score.", "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": "", "extra_scores": {}}, {"review_id": "4dg7TynKJp", "reviewer": "Reviewer_d33o", "summary": "This paper proposes a Symmetry-aware Generative Model (SGM), aiming to learn (approximate) symmetry presented in a data. The model achieves this by mapping each sample onto a prototype—a unique representative on the group orbit—and learning the conditional distribution over its group orbit through maximum likelihood estimation.", "questions": "Please see the above section", "limitations": "Yes", "rating": 4, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is easy to follow", "weaknesses": "1. Initially, I thought the paper aimed to address data-efficient learning of distributions with unknown approximate group symmetry. However, the scope of the paper is quite limited. Essentially, it only aims to learn the conditional distribution of symmetry transformations over the group orbits. In other words, the paper's primary focus is on augmenting training samples over their group orbits according to this supposedly accurate conditional distribution.\n\n2. To demonstrate that the model is functioning as intended, the paper should provide results showing whether this conditional distribution is learned correctly. Unfortunately, this is not shown in any of the examples. The \"realistic-looking\" generated samples depicted in the figures are merely group-transformed versions of the given samples, which is why they appear realistic.\n\n3. One of the main components of the paper is projecting each sample onto a unique representative of its group orbit, achieved through a so-called transformation inference function \\( f_w: X \\to H \\), where \\( H \\) is the (potentially large) group. The paper trains this function \\( f_w \\) in a self-supervised manner through equation (6) to produce unique prototypes. However, this approach is completely unnecessary, as any equivariant function \\( f_w: X \\to H \\) should already accomplish this.\n\n4. The authors claim that one advantage of their approach is handling data sets with unknown symmetries. However, the paper only deals with (2D) rotation and scaling, as well as color transformations, which limits its generalizability.\n\n5. Additionally, the images used in the paper, such as MNIST, are always compactly supported and vanish at the boundary. This makes data augmentation using rotation and scaling feasible. It is unclear how the proposed method would fare when dealing with realistic images that exhibit boundary effects after symmetry transformations.\n\n6. The combination of the proposed model (which supposedly \"learns\" how to augment data) with a VAE is also unconvincing. While it might be better than directly applying data augmentation (as small images can become even smaller), without demonstrating that the conditional distribution is learned accurately, this combined model could very well learn an incorrectly symmetrized distribution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "(1) Learning this conditional distribution is non-trivial. We show that a GAN-based method fails (see Appendix F.1). Furthermore, in sections 3.1 and 3.2 we discuss design and implementation pitfalls that make this challenging. We provide guidelines as to how to resolve them. We consider these sections to be a core part of our contributions.\n\nOur method has several use cases, of which we focus on two:\n* Learning distributions over transformations.\n* Leveraging those transformations to improve deep generative models.\n\nOur experiments demonstrate success in both 1 and 2.\n\n(2) For evidence that our model is functioning correctly see section 4.1 and Appendix F.2. For MNIST, we don’t know the ground truth distribution over the transformations (though our method provides sensible results). However, figures 19 and 20 in Appendix F.2, show results for colored MNIST and dSprites, where we control the transformations. E.g., for colored MNIST, our model learns uniform distributions of hue transformations in exactly the range we added to the dataset. The same is true for dSprites.\n\n(3) Yes, any equivariant function could work, but in general it isn’t obvious how to construct such a function (e.g., a function with equivariance from images to transformation parameterizations for HSV/general affine transformations). Thus, we provide a *general* method for learning such equivariances. In “Invariance of f_ω and the prototypes” we mention the possibility of using an architecture with a subset of equivariances directly built in. We will further clarify this in the camera-ready version.\n\n(4) While we used affine and color transformations in our experiments, our method is general to any transformation for which (approximate) composition and inverse operations exist. This is a very general class of transformations.\n\nWe have shown that our method can successfully learn *3* affine transformations (rotation, shift, scale) and *3* color transformations (hue, saturation, value). These transformations have very different properties. Furthermore, all of these transformations can be learnt concurrently (for a total of 8 transformation parameters).\n\nThis goes beyond several related works which often focus only on affine transformations, and when color transformations are considered it is usually in isolation. E.g.:\n* “Disentangling images with Lie group transformations and sparse coding” by Chau et al.  (from Reviewer TzHS), van der Ouderaa and van der Wilk [2022], Immer et al. [2022] – affine only, \n* Benton et al. [2020], Keller and Welling [2021] – affine and color separately.\n\n(5) Boundary effects may be a problem, by providing a trivial solution for the inference net, allowing it to infer the transformation parameters by ignoring the contents of the image and instead focusing only on the boundaries. \n\nOur results for GalaxyMNIST–where images do not vanish at the boundaries due to other galaxies, stars, and background noise–demonstrate that these boundary effects do not *necessarily* pose a problem for our model as shown in Figures 8d and 12.\n\nNote, there are many transformations (e.g., hue, saturation, and value) for which this isn’t an issue.\n\nNonetheless, we acknowledge this potential limitation, and we'll include a discussion in the camera-ready. This is also a limitation of some existing methods (e.g., LieGAN [Yang et al., 2023], where the edge effects could be used by the discriminator to easily distinguish between real and generated data). It wasn’t the goal of our work to address this limitation.\n\n(6) In addition to the evidence provided in our response to question 2, the fact that our SGM-VAE hybrid models outperform a standard VAE baseline (measured by *test-set* IWLB) shows that the model has learnt a better distribution over transformations than the vanilla VAE. If an incorrect distribution were learned, these models would place too much probability mass on uncommon transformations (and vice versa) which would negatively impact the marginal likelihood of the test data.\n\n> Rating: 2: Strong Reject: For instance, a paper with major technical flaws, and/or poor evaluation, limited impact, poor reproducibility and mostly unaddressed ethical considerations.\n\nWe are surprised that we received this score of 2, which we do not believe reflects our work's quality and contributions. This is especially surprising, given the other reviewers' positive scores and comments, which contradict this score. We hope you will reconsider this score. If you are still convinced that the score is accurate, we kindly ask that you address the following:\n* List the flaws in the paper that you consider major,\n* Explain why you think our evaluation is poor (in contrast with reviewer rf86 who said our paper has “thorough experimental validation on a wide range of datasets” and reviewer GGon who said our paper’s “experiments show interpretable results supporting the claims”), by providing examples of similar papers who’s evaluation we should aim to emulate, \n* Explain why the impact of the paper will be limited (in contrast with reviewers GGon and rf86 who noted the novelty of our work), by citing works that detract from our novelty, \n* Let us know what we can do to improve our reproducibility, given that we provide code and an extensive appendix explaining our experimental setup (“extensive details on the experimental set-ups making them highly reproducible” according to reviewer rf86), or\n* Point to any ethical issues in our paper.\n\nThe review also lists “easy to follow” as the only strength of the paper, which we found surprising given the strengths noted by all other reviewers and the contributions the paper makes upon prior work, which are:\n* A novel generative model (SGM) of the symmetry transformations,\n* A learning algorithm for our SGM,\n* The intuition behind and practical tips for our SGM, and\n* The extensive experimental results that show our model can accurately learn prototypes and distributions over transformations.", "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": "", "extra_scores": {}}, {"review_id": "LxSqkxd3BL", "reviewer": "Reviewer_GGon", "summary": "The paper aims to model the data distribution along each group action orbit. The proposed two-stage method first uses a self-supervised loss to learn an invariant function that maps each sample to its prototype and then uses a normalizing flow to learn the distribution along each orbit (i.e. conditional distribution on each prototype). To acquire the full data distribution, one trains a generative model on the prototypes and composes it with the conditional model. The method shows merits in terms of modeling the prototype-dependent distributions of symmetry transformations in the experiments on small-scale datasets.", "questions": "* Regarding the experiment setting, currently the MNIST dataset is manually transformed. I'd expect smaller variations in the original dataset. In that case, would the proposed method result in less performance increase?\n\n## References\n\n[1] Benton, Gregory, et al. \"Learning invariances in neural networks from training data.\" Advances in neural information processing systems 33 (2020): 17605-17616.\n\n[2] Falorsi, Luca, et al. \"Reparameterizing distributions on lie groups.\" The 22nd International Conference on Artificial Intelligence and Statistics. PMLR, 2019.", "limitations": "The authors have discussed the limitations.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper presents a novel framework for learning the distribution of symmetry transformations. Using a flow generative model allows for much more expressivity, compared to the parametrized Gaussian or uniform distributions considered in [1].\n* The proposed method is easy to understand. The design choices are well-motivated and clearly described.\n* Experiments show interpretable results supporting the claims that\n    * The proposed training objective can lead to approximately invariant prototypes.\n    * The orbital distributions may vary for different prototypes.\n    * The proposed method can increase the data efficiency for datasets with certain degrees of symmetry transformations.", "weaknesses": "* My main concern is the practical applications of the method. Currently, the experiments are done on small image datasets, e.g. dSprites and MNIST. Can the authors identify some more complicated tasks where modeling the symmetry transformations could be beneficial?\n* Also, the dimensionality of symmetry transformations is generally much smaller than the dimensionality of the data manifold. For example, an image dataset may have a 28*28 pixel space and a lot of possible variations in there, while the symmetry of planar affine transformations only accounts for 6 dimensions of variations. Generative modeling is difficult because of the high dimensionality. I'm unsure if it's worth the effort to use a generative model to learn the low-dimensional distribution on a group orbit.\n* (This may be just my personal preference) I find the rotated and colored characters throughout the text a bit distracting. The normal texts have already made things pretty clear. Those characters may be great for intuition but are less accurate and formal.\n* Currently, the parameterizations for different symmetry transformations seem ad-hoc. E.g. affine transformations are represented in the affine matrix. However, symmetry groups have different structures (which can, for example, be reflected by the structure constants of the Lie algebra) and require different ways of parametrizing distributions. The authors should address this aspect and possibly discuss some related works, e.g. [2].", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> My main concern is the practical applications of the method. Currently, the experiments are done on small image datasets, e.g. dSprites and MNIST. Can the authors identify some more complicated tasks where modeling the symmetry transformations could be beneficial?\n\nPlease see our general response, in which we provide detailed motivation as to why our experiments are sufficiently interesting and informative. In short, we have already gone beyond many published related works, our choices of datasets cover a range of dataset sizes and data dimensionalities, and the GalaxyMNIST dataset contains high-dimensional natural images while providing challenges due to its small number of observations and large number of symmetries.\n\n> Also, the dimensionality of symmetry transformations is generally much smaller than the dimensionality of the data manifold. For example, an image dataset may have a 28*28 pixel space and a lot of possible variations in there, while the symmetry of planar affine transformations only accounts for 6 dimensions of variations. Generative modeling is difficult because of the high dimensionality. I'm unsure if it's worth the effort to use a generative model to learn the low-dimensional distribution on a group orbit.\n\nWe agree that the data manifold is larger than the dimensionality of the transformations and that one of the main challenges of generative modeling is high dimensionality. This is actually one of the main motivations for this work. Our hypothesis is that decomposing the generative modeling task into an “easier” symmetry modeling task and a more complicated task of modeling all of the other sources of variation, will provide benefits such as data-efficiency and improved model fit. Figures 11 and 12 confirm this hypothesis, with our symmetry augmented VAEs easily outperforming vanilla VAEs without this inductive bias. In addition to these improvements, this decomposition provides benefits in interpretability of the latent code, and the ability to easily generate realistic “data-augmentations” of any observation. Whether or not it is worth the effort depends on the particular application at hand, however, we are confident that the strengths of our method could be useful in several settings. Finally, we hope that our “novel framework” and “easy to understand” method will result in further interesting developments within the field.\n\n> (This may be just my personal preference) I find the rotated and colored characters throughout the text a bit distracting. The normal texts have already made things pretty clear. Those characters may be great for intuition but are less accurate and formal.\n\nWe appreciate your input on this, and understand that their inclusion might not be to everyone’s tastes. However, prior to their inclusion, we received several pieces of feedback that the corresponding sections were hard to understand. Since their inclusion we have found that readers are much more easily able to understand the text. Furthermore, yourself as well as reviewers d330 and rf86 have noted the clarity of our text. Nonetheless, we will take your input to heart, and for the camera-ready version we will carefully consider the inclusion of each of these characters.  \n\n> Currently, the parameterizations for different symmetry transformations seem ad-hoc. E.g. affine transformations are represented in the affine matrix. However, symmetry groups have different structures (which can, for example, be reflected by the structure constants of the Lie algebra) and require different ways of parametrizing distributions. The authors should address this aspect and possibly discuss some related works, e.g. [2].\n\nWe will happily include a discussion of [2] (and any other related works you suggest). However, we note that we use both affine matrices and the corresponding Lie algebra constants when representing affine transformations, as each representation has pros and cons and the two representations can easily be interchanged. For instance, our inference and generative networks output the Lie algebra constants, since these are low dimensional and easy to constrain (if necessary). On the other hand, we use affine matrices for transformation composition. Both of these choices are discussed in Section 3.1. Similarly, our representation of color transformations is chosen to make composition of transformations simple. See Appendix D.6 for further details. \n\n> Regarding the experiment setting, currently the MNIST dataset is manually transformed. I'd expect smaller variations in the original dataset. In that case, would the proposed method result in less performance increase?\n\nYour intuition is correct – the smaller the degree of transformation present in the dataset, the less our method is expected to improve data efficiency. This can be observed in Figure 11 – we see that as more rotation is added to the dataset, the performance gap between AugVAE and VAE becomes larger. However, our results of GalaxyMNIST (which don’t include any manual transformations) demonstrate that our method still performs well in natural settings. Finally, we note that the performance gain from our method can be increased by including a wider range of transformations (e.g., learn both color and affine transformations together).", "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": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their time and constructive feedback on our paper. We are pleased that the reviewers have highlighted the quality of our writing (rf86, d33o, GGon), experimental evaluation (rf86, TxHS, GGon), the novelty, elegance and interestingness of our work (rf86, TxHS, GGon). We would like to highlight the following quotes from the reviews:
* “The paper is very well written and proposes an elegant method”, “The paper is very well-written and has beautiful illustrations”, “The paper provides very thorough experimental validation on a wide range of datasets”, “The appendices provide extensive details on the experimental set-ups making them highly reproducible” – reviewer rf86
* “The overall architecture seems interesting and sound”, “the idea of learning probability distribution over the group elements is especially interesting” – reviewer TzHS
* “The paper is easy to follow” – reviewer d33o
* “The paper presents a novel framework for learning the distribution of symmetry transformations.”, “The proposed method is easy to understand.”, “The design choices are well-motivated and clearly described.”, “Experiments show interpretable results supporting the claims …” – reviewer GGon

We also acknowledge a common criticism among most of the reviewers (TzHS, d33o, GGon) that the paper would be improved with the addition of a larger-scale dataset with ‘natural’ images. We believe our experiments are sufficiently interesting and informative for the following reasons.

Many *published* related works do not go beyond image datasets of similar size/dimensionality to MNIST and dSprites. Examples of this include:
* “Disentangling images with Lie group transformations and sparse coding” by Chau et al., and “Group Equivariant Subsampling” by Xu et al. (mentioned by Reviewer TzHS), and
* Yang et al. [2023], Benton et al. [2020], van der Ouderaa and van der Wilk [2022], Immer et al. [2022], Kaba et al. [2023], Keller and Welling [2021], and Bouchacourt et al. [2021a], from our related work
to name just a few. 

Unfortunately, it is impractical for us to go beyond these settings due to computational resource limitations.

Our set of experiments using dSprites, MNIST, and GalaxyMNIST, considering different types of symmetry transformations on each, demonstrate the general applicability of our method from small to large data sizes, from small to large dimensionalities, and for several different symmetries. The dSprites images, while simple, have a fairly large dimensionality (64 x 64 pixels) and are very plentiful (~740k images). GalaxyMNIST contains a small number of images (only 7k for training) of even larger dimensionality (64 x 64 x 3). This small data regime is perhaps more interesting since it demonstrates that our method is able to accurately capture the distribution over transformations without much data. Furthermore, we view the GalaxyMNIST images as ‘natural’ in that they come from real-world astronomy observations collected for the Galaxy Zoo DECaLS Campaign. Finally, these datasets demonstrate that our model can learn five affine transformations alone (MNIST, dSprites), three color transformations alone (MNIST), and both affine and color transformations together (GalaxyMNIST). 

Like those excellent papers mentioned above, this paper has focused on novel ideas and methodological contributions rather than large-scale experiments. While we also value those experiments, we note that such engineering work often follows from work that builds up understanding and provides a promising proof-of-concept. Furthermore, we note that scaling up generative models is typically harder than their deterministic counterparts. Thus, we hope this will not be considered a major weakness of our work.

Meta-review: This is a well-written and well-motivated paper that presents a generative model of symmetry transformations. The main idea uses some recent work that treat relaxed symmetry as a latent variable over which inference can be performed. The experiments are thorough and well-described and validate the approach sufficiently. The reviews for the paper were mostly positive with one negative review. However, during the discussion period, the authors were partly able to address some of the reviewer concerns. On reading the paper, I tend to agree that the paper makes a good contribution and could be a valuable addition to the literature. The authors are requested to use the extra page to incorporate various points that have come up during the discussion phase.
Decision: Accept (poster)