A Generative Model of Symmetry Transformations

Authors: james urquhart allingham
bruno kacper mlodozeniec
shreyas  padhy
javier  antorán
david  krueger
richard e turner
eric  nalisnick
josé miguel hernández-lobato
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-03 16:50:06

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.

Main paper: 
 learning: 
 We now discuss learning for the two NNs required by our model, f ω (x) and p ψ (η | x). In Appendix A, we connect our learning algorithm with MLL optimization using an ELBO.For T -1 η , with η given by f ω , 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 T -1 fω(x) (x) should be invariant to transformationsEQUATIONTo learn such a function, we optimize for this property directly. To this end, we sample transformation parameters η rnd from some distribution over parameters p(η rnd ). This allows us to get random samples x rnd := T ηrnd (x) ∈ X in the orbit of any given element x ∈ X . Since we want full (i.e., strict) invariance, p (η rnd ) must have support on the entire orbit [van der Ouderaa and van der Wilk, 2022]. We then learn an equivariant via a self-supervised learning (SSL) schemef ω3 inspired by methods like BYOL [Grill et al., 2020] and, more directly, BINCE [Dubois et al., 2021] . For example, we could use the objective illustrated in Figure 4 :EQUATIONOur actual objective differs slightly. SinceEQUATION)This change allows us to reduce the number of small discretization errors introduced with each transformation application by replacing repeated transformations with a single composed transformation; see Section 3.1 for further discussion. Our SSL loss is given in line 1 of Algorithm 1.Require: initial parameters ω init & ψ init , dataset D 1: function SSL_LOSS(x, ω)2:η x ← f ω (x) 3: η rnd ∼ p (η rnd ) 4:x rnd ← T ηrnd (x)5:η xrnd ← f ω (x rnd ) 6: x ′ ← T ηx • T -1 ηx rnd (x rnd ) 7: output mse(x, x ′ ) 8: function MLE_LOSS(x, ω, ψ) 9: η x ← f ω (x) 10: x ← T -1 ηx (x) 11: output -log p ψ (η x | x) 12: ω, ψ ← ω init , ψ init 13: while ω not converged do 14: X ← next_batch(D) 15: update ω with ∇ ω 1 B B b=1 SSL_LOSS(X b , ω) 16: while ψ not converged do 17: X ← next_batch(D) 18: update ψ with ∇ ψ 1 B b MLE_LOSS(X b , ω, ψ) 19: output ω, ψGenerative model of transformations. Once we have a prototype inference function, we simply learn p ψ (η | x) by maximum likelihood on the created data pairsf ω (x i ), T -1 fω(xi) (x i ). This is shown in line 8 of Algorithm 1. While we need to specify the kinds of symmetry transformations T η we expect to see in the data, by learning p ψ (η | x) the model can learn the degree to which those transformations are present in the data. Thus, we can specify several potential symmetry transformations and learn that some are absent in the data. Furthermore, the required prior knowledge (the support of p (η rnd )) is small compared to what our SGM can learn (the shapes of the distributions for each of the present transformations).Since we are primarily interested in using the model to (a) inspect the distribution over naturally occurring transformations for a given element x, and (b) resample new "naturally" augmented versions of the element, we do not need to learn p (x). We can do (a) by querying p (η | x = x) for x := T -1 fη(x) (x), and we can do (b) by sampling η ∼ p (η | x) and transforming the x to get x := T η ( x). Of course, if one wanted to sample new prototypes, one could fit p θ (x) using, e.g., a VAE. Not learning p (x) greatly simplifies training for complicated datasets that would otherwise require a large generative model, an observation made by Dubois et al. [2021] . x x p (η | x, x) 8 8 0.5 • δ(η -0 • ) + 0.5 • δ(η -180 • ) 8 8 0.5 • δ(η -30 • ) + 0.5 • δ(η + 150 • ) 8 8 0.5 • δ(η + 30 • ) + 0.5 • δ(η -150 • ) 8 8 8 8 8 8 0 • η | 8 (b) Simple p ψ (η | x) 8 8 8 8 8 8 0 • η | 8 (c) Flexible p ψ (η | x) 
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.Working with transformations. Repeated application of transformations-e.g., in Figure 4 -can 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 affine-transformation matrices. More generally, if there is some representation of the transformation parameters T (η) where composition can be performed-e.g., as matrix multiplicationT η2 • T η1 = T ′ T (η2)T (η1), in the case where T is a group representation-then 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 a 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 out-of-bounds4 ). If appropriate bounds are known a priori, we can prevent severe information loss by constraining η min and η max using tanh, scale, and shift bijectors. Alternatively, we can augment the SSL loss in Algorithm 1 with an invertibility lossEQUATION)Learning p ψ (η | x) with imperfect inference. In practice, our transformation inference network f ω (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 p ψ (η x | x) robust to these variations, we train it with prototypes corresponding to randomly transformed training data points. I.e., we modify the MLE objective in Algorithm 1 as log p ψ (η x | x′ ), where x′ = T -1 fω(Tη rnd (x)) (T ηrnd (x)) as in our SSL objective. Averaging the loss over multiple samples-e.g., 5-of η rnd is beneficial.We now motivate some of the design choices for our SGM by means of illustrative examples. In each case, we assume that T η is counter-clockwise rotation; thus, η is the angle.1. The distribution p ψ (η | x) is implemented as a normalising flow. Consider a dataset of '8's rotated in the range -30 • to 30 • : { 8, . . . , 8, . . . , 8 }. Let us assume that the prototype is '8'. Figure 5a shows p (η | x, x), an example of the true distribution for η given x and x, for several observations, under the data generating process5 . These distributions are composed of deltas because 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-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.x x p (η | x, x) 2 2 δ(η -0 • ) 2 2 δ(η -30 • ) 2 2 δ(η + 30 • ) 8 8 0.5 • δ(η -0 • ) + 0.5 • δ(η -180 • ) 8 8 0.5 • δ(η -30 • ) + 0.5 • δ(η + 150 • ) 8 8 0.5 • δ(η + 30 • ) + 0.5 • δ(η -150 • ) η (b) η η | 2 η | 8 (c) η | x2. The transformation parameters η depend on the prototype x. Consider a dataset of '2's and '8's rotated in the range -30 3. The prototype x is fully invariant to transformations of x. Models such as CNNs are most useful when we know a priori which symmetries 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 a model that is invariant to rotations in the full range [-180 • , 180 • ] might be unable to distinguish between '6' and '9'. Thus, in the literature for learning invariances in the discriminative setting, it is common to learn partially invariant functions that capture some degree of invariance [van der Wilk et 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 x contains no information about the symmetries in the data.x x p (η | x, x) x p (η | x, x) x p (η | x, x) 2 2 δ(η -0 • ) 2 δ(η + 15 • ) 2 δ(η -0 • ) 2 2 δ(η -30 • ) 2 δ(η -15 • ) 2 δ(η -0 • ) 2 2 δ(η + 30 • ) 2 δ(η -0 • ) 2 δ(η -0 • ) η | 2 (b) FULL η | 2 η | 2 (c) PARTIAL η | • (d) NONEConsider a dataset of '2's rotated in the range -30 • to 30 • : { 2, . . . , 2, . . . , 2 }. Figure 7a shows predicted prototypes and the corresponding distributions over η for several observations. There are three cases: (a) a fully-invariant x, i.e., there is a 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 partiallyinvariant case to the extreme and has as many prototypes as observations. In the partially-invariant and non-invariant cases, we can get multiple prototypes rather than a single unique prototype per orbit, which is invalid under the generative model of the data. As a result, p ψ (η | x) does not represent the distribution of naturally occurring transformations of x in the data. This is illustrated in Figures 7b to 7d, which show idealized examples of the learned p ψ (η | x) in 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 is a problem, let us say we would like to probe the probability of a particular transformed variant of an observed example. For example, given an example of a digit '3', we want to know the probability of observing , that digit rotated by -90 • . Assuming we can find a prototype x we would like p (η | x = x) to represent all naturally occurring augmentations. Unless x is unique, this won't necessarily be the case, as illustrated in Figure 7 .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 crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
experiments: 
 In Section 4.1, we explore our SGM's ability to 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, we leverage our SGM to improve data efficiency in deep generative models.We conduct experiments using three datasets-dSprites [Matthey et al., 2017] , MNIST, and GalaxyM-NIST [Walmsley et al., 2022] -and two kinds of transformations-affine and color. In Section 4.1, when working with MNIST under affine transformations, we add a small amount of rotation (in the range[-15 • , 15 • ]) to the original data to make rotations in the figures easier to see. For MNIST under color transformations, we first convert the grey-scale images to color images using only the red channel. We then add a random hue rotation in the range [0, 0.6π] and a 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 case of the heart sprites, we have removed the rotations and set the y-positions to be bimodal in the top and bottom of the images. Further details about the dSprites setup, as well as 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 a group that is easy to work with. We also learn color transformations (hue, saturation, and value). See Appendix C.7 for details about how we parameterize T η in both cases.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 realistic "natural" examples that are almost indistinguishable from test examples. There are sev-eral illustrative examples which bear further discussion. The heart sprites in Figure 8a show that our SGM was able to learn the absence of a 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 in 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 distributions (on rotation in this case). The figure's third '7' is interesting because our SGM interprets it as a '2'.-1 2 1 2 -1 2 1 2 -π π -1 1 -1 1Flexibility is important. In η, each dimension corresponds to a different transformation. We refer to p ψ (η i | x) as the marginal distribution of a single transformation parameter. Figure 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-the 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 a much smaller 2 nd 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 p ψ (η | x), as conjectured in Section 3.2. We also find significant dependencies between dimensions of η (e.g., rotation and translation in dSprites). Invariance of f ω 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 transformation parameters at each iteration, normalized by the maximum magnitude (at iteration 0). The first prototype x1 is most different from the previous x0 = x, with successive prototypes being similar visually and as measured by the magnitude of the inferred transformation parameters. However, the magnitude of the inferred parameters does not tend towards 0, rather plateauing at around 5% of the maximum. This highlights that, although simple NNs can learn to be approximately invariant, a natively invariant architecture has the potential to improve performance.We use SGM to build data-efficient and robust generative models. In Figure 11 , we compare a standard VAE to two VAE-SGM hybrid models-"AugVAE" and "InvVAE"-for 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 ±θ max , the maximum added rotation angle). Thus, adding rotation here is not data augmentation. AugVAE is a VAE that uses our SGM to re-sample transformed examples x ′ = T η| x ( x), introducing data augmentation at training time. InvVAE is a VAE that uses our SGM to convert each example x to its prototype x at both train and test time. That is, the VAE in InvVAE sees only the invariant representation of each example. We also compare against a VAE trained with standard data augmentation6 . We use test-set importance-weighted lower bound (IWLB) [Domke and Sheldon, 2018] As expected, for the VAE ( ), as we decrease the amount of training data ( → ) or increase the amount of randomly added rotation, performance degrades. This is because the VAE sees fewer training examples per-degree of rotation. On the other hand, the AugVAE ( ) is more data efficient. Its performance is unaffected by reducing the number of observations by three quarters. Furthermore, while the performance of AugVAE and the standard VAE are almost identical for small angles and large training sets, the drop in performance of AugVAE for larger random rotations is significantly smaller; AugVAE does not see less training examples per-degree of rotation. InvVAE ( ), 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 ( ) improves on the standard VAE for less training data, it is substantially worse in the presence of more data. This contrasts our AugVAEs, 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 in accordance with the true data distribution. We further validate these results with the more complex GalaxyM-NIST dataset and an enlarged set of both affine and color transformations. As with our rotated MNIST with affine transformation results, in Figure 12 [2023] show that such canonicalization functions can be used to make large-pre-trained NNs equivariant and, when combined with dataset-dependent symmetry priors, do not degrade performance. Similarly, Kim et al. [2023] learn architecture-agnostic equivariant functions by averaging a non-equivariant function over a probabilistic prototypical 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] .Data augmentations and symmetries. Prior work makes several connections between data augmentation 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 our approach of learning data augmentation from the data and our architecture-agnostic 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-dependent.Symmetry-aware latent spaces. Encoding symmetries in latent space is well-studied. Higgins et al. [2018] posit that symmetry transformations that leave some parts of 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 model-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] 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] 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, Aliee 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] construct a VAE with a topographically organised latent space such that an approximate equivariance is learned from sequences of observations. In contrast to 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 a 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 to disentangle sources of variation in a dataset, thereby learning a representation that is equivariant to each of the sources while invariant to all others.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?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-including dataset splits, model architectures, hyper-parameter settings and sweeps, transformation parameterisations, and a 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 a 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 a specific model and empirical evaluation, it may be necessary to either make it 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., 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 submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example (a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration 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 a direct path to negative applications. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment 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 to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an 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 to 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 or 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 a system learns from feedback over time, improving the efficiency and accessibility of ML). 
conclusion: 
 We have presented a Symmetry-aware Generative Model (SGM) and demonstrated that it is able to learn, in an unsupervised manner, a distribution over symmetries present in a dataset. This is done by modeling the observations as a random transformation of an invariant latent prototype. This is 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 in 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 et 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 be 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. 
safeguards: 
 Question: Does the paper describe 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, or scraped datasets)? 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 we encourage authors to take this into account and make a 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?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 an 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 or 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 we 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 as the institution conducting the review. 
: 
 generative process that we aim to model in this paper. A prototype x ( ) is transformed by Tη into an observation x ( , , ). The transformatione.g., rotation-is parameterized by η-e.g., an angle. Right: The corresponding orbit-i.e., the set of all possible instances of x that can result from applying Tη-with a few elements shown. Under this generative process, the prototype is an arbitrary orbit element. Each element in the orbit has a probability p (x | x) induced by p (η | 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.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 symmetries [LeCun et al., 1989 ]-and can be extended to reflection and rotation [Cohen and Welling, 2016] -and more recently, transformers, with permutation symmetries [Lee et al., 2019] .In many cases, it is not known a 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 et al., 2023 , Mlodozeniec et al., 2023] . However, in these works-which focus on invariant discriminative 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, a generative model must capture the factors of variation corresponding to the symmetry transformations of the data. Doing so can provide benefits such as 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 a 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 a given example x, and resample new "naturally" augmented versions of the example. Our contributions are 1. We propose a Symmetry-aware Generative Model (SGM). The SGM's latent representation is separated into an invariant component x and an equivariant component η. The latter, η, captures the symmetries in the data, while x captures none. We recover x by applying a parameterised transformation, x = T η (x). We call x a prototype since each x can produce arbitrarily transformed observations; see Figure 1 . 2. We propose a two-stage algorithm for learning our SGM: first learning x using a selfsupervised approach and then learning η via maximum likelihood. Importantly, this does not require modeling the distribution of prototypes p (x), 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-likelihood can improved by using our SGM to incorporate symmetries. Additionally, unlike a standard VAE, explicitly modeling symmetries makes our VAE-SGM hybrid robust to deleting half of the dataset.Notation. We use a, a, and A (i.e., lower, bold lower, and bold upper case) for scalars, vectors, and matrices, respectively. We distinguish between random variables such as x, η, A, and their realizations x, η, A. Thus, for continuous a, p (a) is a PDF that returns a density p (a = a) = p (a).We use • to represent function composition, e.g.,f 1 • f 2 .2 Symmetry-aware Generative Model (SGM)Consider a dataset of observations {x n } N n=1 on a space X , and a collection {T η } of transformations T η : X → X parameterised by transformation parameters η ∈ H ⊆ R dη . We assume {T η } η∈H (abbreviated {T η }) form a group. Loosely, our aim is to model the distribution over transformations present in the data. To do so, we model the distribution p (x) by decomposing it into two disparate parts: (1) a distribution over prototypes and (2) a distribution over parameters controlling transformations to be applied to a prototype. Concretely, we specify our generative model as follows (also depicted in Figure 2 ):x ∼ p (x) , (1) η ∼ p ψ (η | x) ,(2) x = T η (x).(3)T η (x) p (η rnd ) η rnd -50 • f ω (x) x rnd 5 T -1 η (x) η xrnd -80 • mse (•, •) x′ 5 f ω (x) T -1 η (x) η x -30 • x 5 L SSLFigure 4 : Self-supervised symmetry learning. We encourage f ω (x) to be equivariant by mapping x and a randomly transformed x to the same x. Gray text shows examples for each variable in the graph. Note that x and x rnd may not appear in the dataset; see Figure 1 .That is, the SGM assumes that each observation x is generated by applying a transformation T ηparameterized by a latent variable η-to a latent prototype x. Since x, by assumption, contains no information about the symmetries in the data, p ψ (η | x) must model the distribution over the transformations T η present in the data. Motivation. Why would we expect specifying p (x) in this way to be useful? Firstly, our SGM allows us to query a distribution over naturally occurring transformations p ψ (η | x = T -1 η (x)) for any input x, given the matching prototype x := T -1 η (x). Secondly, we expect our SGM to align with the true physical process of generating the data for many interesting datasets. As an illustrative example, when a person writes a digit, they first decide what kind of digit to write-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 1 . Similarly, when a photographer captures an object, the photo is also a function of latent factors of variation, such as lighting, the lens, camera shake, etc.What do we require of a prototype? x can informally be considered a canonical/reference example with no transformation applied to it. More precisely, we require that for any orbit of an element x-defined as the set of elements in X which x can be mapped to by a transformation in {T η }-there is exactly one prototype in the orbit. Figure 1 depicts an example orbit-a set { , , , ...} of all rotated variants of a '3'-with a unique prototype. (x 1 , x 2 ) is transformed via T η : (x 1 , x 2 ) → (x 1 , x 2 )+(η, 0).Thus, horizontal lines form disjoint orbits in 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 T η .Why do we want a group? Having the transformations {T η } be a group simplifies things, since {T η } will then naturally partition the space X into (disjoint) orbits. Within each orbit, every element can be transformed into one another with a transformation in {T η }. As an example of such a partition, if our collection of transformations were horizontal shifts T η : x → x + (η, 0) acting on a point x ∈ R 2 , then the different orbits will correspond to all points on a given horizontal line; see Figure 3 . Therefore, if we have chosen a unique prototype for each orbit and {T η } forms a group, any two elements x, x ′ ∈ X will have the same prototype if and only if they can be transformed into one another.In Section 2.1, we describe a method for learning a transformation inference function f ω : X → H, with parameters ω, that for x ∈ X returns transformation parameters η ∈ H as η = f ω (x). These map x to a prototype x := T -1 η (x) that generates x := T η (x) 2 . We then apply standard generative modeling tools to learn p (x, η) = p (x) p ψ (η | x) given the generated data pairs { xn , η n } N n=1 .1 Our SGM does not always perfectly match the data-generating process. E.g., a person is unlikely to "imagine" the same prototype for both a '6' or a '9'-which can often be transformed into one another with rotation. 2 The transformation is not necessarily unique. 
 
 Appendix: As we will now show, Algorithm 1 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:EQUATIONwhere p θ (x) is some generative model-e.g., a VAE-for prototypes, with parameters θ, andq ω (η, x | x) = q ω (η | x) p (x | x, η). Now, we can show that the gradient of the likelihood term in the ELBO is approximated by the gradient of our SSL loss on line 1 of Algorithm 1:EQUATIONEQUATION▷ take 1 sample, η ∼ q ω (η | x):EQUATION) ▷ definition of Gaussian PDF:EQUATION▷ drop constant term:EQUATION)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. ψ is approximated by the gradient of our MLE loss on line 8 of Algorithm 1:EQUATION▷ drop constant terms and use x = T -1 η (x) :EQUATION▷ take 1 sample, η x ∼ q ω (η | x):EQUATIONNote 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 instance, we do not learn the generative and inference models jointly. This disjoint training is equivalent Prototypes produced by VAE who's encoder is made invariant using ( 22), where η ∼ U (-η max , η max ) and η max = (0.25, 0.25, π, 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∇ ω D KL [q ω (η, x | x) || p ψ (η | x) p θ (x)]when training q ω (η | x). This KL-divergence has two components: entropy -H [q ω ] and cross entropy H [q ω , p ψ p θ ]. Assuming that p ψ (η | x) is sufficiently flexible, the cross entropy term should not have a significant impact on q ω (η | x) since p ψ is trained to match q ω . On the other hand, q ω (η | x) should be close to a delta since there should be a single prototype for each x. Thus, encouraging high variance with an entropy term might actually be harmful. Another difference is that we do not need to learn p θ (x), which has the benefit that we can learn the symmetries in a 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 to the fact that p (x | x, η) is a 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 p ω (η | x) as being closely related to the process of learning a prior in an ELBO.In an early version of this work [Allingham et al., 2022] , we trained a variant of the SGM using an ELBO similar to (12), with the main difference being that x was modeled using a VAE and invariance was incorporated into the VAE encoder. We constructed an invariant encoder q ϕ (z | x) from a non-invariant encoder qϕ (z | x):EQUATIONfollowing Benton et al. [2020] , van der Ouderaa and van der Wilk [2022], Immer et al. [2022] . We found that this approach worked well for a 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.This section elaborates on Section 3.1 and provides additional considerations.Suitability of NN architectures. The architecture of f ω must be compatible with learning an equivariant mapping from x to η. For example, a standard CNN requires many convolutional filters to represent a function that is (approximately) equivariant to continuous rotations [Maile et al., 2023] .X -space vs. H-space SSL objective. One might notice that it is possible to remove the T -1 η 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 -space version is preferred since different parameters η 1 , η 2 can map to the same transformed element T η1 (x) = T η2 (x). E.g., consider rotations transformations applied to various shapes: for a square T 0 • ≡ T 90 • ≡ T 180 • ≡ T 270 • all map to the same transformed image, and an H-space objective incorrectly penalizes differences of ±n × 90 • in η values.We compare rotation inference 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 -4 and a batch size of 256-trained on fully rotated MNIST digits using both X -space and H-space SSL objectives: Objective x-mse η-mse X -space 0.2387 0.9715 H-space 0.3567 0.4736 average of X -space and H-space 0.3129 0.4619 When using the H-space objective, we see the distance in observation (X ) space.Learning q ω (η|x) instead of f ω . We found that learning f ω probabilistically-i.e., allowing for some uncertainty in the transformation during the training process by parameterizing a density over H with q ω (η|x) and sampling η-provides small improvements in performance. The distribution q ω (η|x) quickly collapses to a delta. Thus, we hypothesize that the added noise from sampling acts as a 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 hyperparameters, training the inference network fails, and the prototype transformations would be 100% lossy-i.e., they would result in completely empty images-regardless 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 to improve performance, so too is 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 AdamW optimizer with a cosine decayed with warmup learning rate schedule that starts at 1 × 10 -4 , increases to 3 × 10 -4 in 500 steps, and then decreases to 1 × 10 -7 , with a batch size of 256-trained on fully rotated MNIST digits using the SSL objective averaged over 1, 3, 5, 10, and 30 samples: As the number of samples increases, x-mse decreases until saturating around 5 samples. Note that this relationship is not likely to be monotonically decreasing because there is random noise in each training run (i.e., due to 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 a 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 x and x rnd (a randomly transformed version of x). An alternative is to compare the prototypes given x rnd1 and x rnd2 , two randomly transformed versions of x:EQUATION)As before, we modify this loss to allow us to compose transformations to getEQUATION)The motivation for using this 'symmetric' SSL loss is that it provides the inference network with additional data augmentation-the inference network is now unlikely ever to see the x twice. We find that while this works well for MNIST, it does not work well for dSprites. This is because the transformations in dSprites in dSprites are more lossy than those for MNIST. E.g., it is easier to shift a small sprite out of the frame of an image compared to a large digit. Thus, the symmetric loss results in a much higher variance when used with dSprites, which negatively impacts training.Composing affine transformations of images. Care must be taken when composing affine transformations of images when implemented via a coordinate transformation (e.g., affine_grid & affine_sample in PyTorch, or scipy.map_coords in Jax). To compose two affine transformations parameterised by η 1 and η 2 , the affine matrices T (η 1 ), T (η 2 ) need to be right-multiplied with one another; in other wordsT η2 • T η1 = T ′ T (η 1 )T (η 2 ). 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-transformed by T -1 .Overfitting of the generative network. While we did not observe any overfitting of the inference network (likely due to 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 p ψ (η | x) with imperfect inference, continued. To encourage p ψ (η | x) produce the same distribution for the inconsistent prototypes produced by q ω (η | x), we add a consistency loss to line 8 of Algorithm 1 the MLE objective:EQUATIONwherep i = p ψ (η x | x′ i )and x′ i is due to the i th η rnd sample.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, ml_collections to specify our configurations, and wandb to track our experiments. The code is available at https://github.com/ cambridge-mlg/sgm.Unless otherwise specified, we use the following NN architectures and other hyperparameters for all of our experiments. We use the AdamW optimizer with weight decay of 1×10 -4 , 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. In some cases, we use Dropout. The structure of each layer is Dense → gelu → LayerNorm → Dropout. Whenever we learn or predict a scale parameter σ, 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 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 a small amount of blur to the images with a Gaussian filter of size 5 for the first 1% of training steps. The σ value for the filter was linearly decayed from their maximum to 0. The initial maximum value is specified below. The dataset consists of sprites with the outer product of these factors, for a total of 737280 examples. We modified our data loader to resample the sprites proportional to the following distributions on the latent factors conditioned on the shape.• square -Scale: TruncNorm µ = 0.75, σ 2 = 0.2, min = 0.55, max = 1.0 -Orientation: U (0.0, 2π) -X position: U (0.5, 0.95) -Y position: U (0.5, 0.95)• ellipse -Scale: TruncNorm (0.65, 0.15, 0.5, 0.85) -Orientation: U (0.0, π/2) -X position: TruncNorm (0.5, 0.25, 0.1, 0.9) -Y position: TruncNorm (0.5, 0.15, 0.35, 0.65)• heart -Scale: U (0.9, 1.0) -Orientation: δ (0.0) -X position: U (0.1, 0.5) -Y position: 0.5 • U (0.1, 0.3) + 0.5 • U (0.7, 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.We make use of the GalaxyMNIST dataset [Walmsley 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 the inference and generative networks, we constrain their outputs to be in [-η max , η max ] + (0., 0., 0., 0., 0., 0.5, 0., 0.), where η max = (0.75, 0.75, π, 0.75, 0.75, 0.5, 2.31, 0.51) using with tanh and scale bijectors. This dataset contains 10k examples. We use the last 2k as our test set, and the previous 1k as a validation set.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 parameters from U (-η max + (0., 0., 0., 0., 0., 0.5, 0., 0.), η max + (0., 0., 0., 0., 0., 0.5, 0., 0.)), where η max 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 GalaxyMNIST results use a blur σ init of 0, a gradient clipping norm of 10, a learning rate of 3×10 -4 , an initial learning rate multiplier of 1×10 -2 , a final learning rate multiplier of 3×10 -4 , and a warm-up steps % of 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.Generative network. We randomly augment the inputs by sampling transformation parameters from U (-η max × 0.75 + (0., 0., 0., 0., 0., 0.5, 0., 0.), η max × 0.75 + (0., 0., 0., 0., 0., 0.5, 0., 0.)), where η max matches the constraints above. For the VAE data-efficiency results in Figure 12 , we perform the same hyperparameter grid search as above for each random seed and amount of training data, with the following changes. 7 The sweep for number of training steps is [3.75k, 7.5k, 15k] . All of the other GalaxyMNIST results use a 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.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 × 96 pixels to 64 × 64 using bilinear interpolation. The dataset has dedicated train, test, and validation splits which we 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 a single random seed.Inference network. We train for 20k steps.Generative network. The sweep for number of training steps is [15k, 30k, 60k].8 Our VAEs use a latent code size of 20. The prior is a normal distribution with learnable mean and scale, initialized to 0s and 1s, respectively.Our VAE encoders are LeNet-style CNNs with convolutional feature extractors followed by an MLP with a 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 3×3 filters. The first two Conv have a stride of 2, while all others have a stride of 1. In between the convolutional layers and the MLP, there is a special dimensionality reduction Conv with only 3 filters followed by a flatten. For each dimension of the latent code, the encoder predicts a mean µ and a scale σ. The means and scales are initialized to 0s 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 are 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 a homoscedastic per-pixel scale σ. The scales are initialized to 1.We use an initial learning rate multiplier of 3×10 -2 , and a final learning rate multiplier of 1×10 -4 . We run the following grid sweep for each (random seed, number of training examples, maximum added rotation angle) triplet:• learning rate ∈ [3×10 -3 , 6×10 -3 , 9×10 -3 ],• convolutional filters ∈ [(64, 128), (64, 128, 256) ],• number of training steps ∈ [5k, 10k, 20k], and • warm-up steps % ∈ [0.15, 0.2].When running the sweep for AugVAE and InvVAE, we use the inference and generative network hyperparameters for the same (random seed, number of training examples) pair.For our PatchCamelyon experiments, we use only a single random seed and a slightly modified hyperparameter sweep: (64, 128), (64, 128, 256), (128, 256, 512) ],• learning rate ∈ [3×10 -3 , 6×10 -3 , • convolutional filters ∈ [• number of dense hidden layers ∈ [1, 2],• number of training steps ∈ [20k, 30k, 40k], and • warm-up steps % ∈ [0.15].We consider five affine transformations: shift in x, shift in y, rotation, scaling in x, and scaling in y. We represent these transformations using affine transformation matricesA = exp ( i η i G i ),where G i are generator matrices for rotation, translation, and scaling; see Benton et al. [2020] . The transformations are applied to an image by transforming the coordinates (x, y) of each pixel, as in Jaderberg et al. [2015] :[x ′ y ′ 1] ⊺ = A • [x y 1] ⊺ .To parameterize color transformations, we use an equivalent representation of color images in Hue-Saturation-Value (HSV) space, where each pixel is represented as a tuple(h, s, v) ∈ {[-π, π] × [0, 1] × [0, 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 vertical axis, the saturation corresponds to the radial distance from the cone's center, and the value corresponds to the distance along the cone's vertical axis, with a value of 0 corresponding to the tip of the cone, and a value of 1 corresponding to the base of the cone. We color-transform an image by transforming each pixel asEQUATIONWe therefore obtain η= (η h , η s , η v ) ∈ {[0, 1] × R × R}.We choose this specific form of parametrizing the η parameters in order to gain the convenience of simply adding and subtracting in η space when carrying out color transform compositions and inverses. More concretely, with our chosen parametrization, we obtain the property that T η1 • T η2 = T η1+η2 . Therefore, we can easily perform compositions and inversions in η space for color transformations without resorting to matrix multiplications. In order to achieve this, we first consider hue, which is easy to parametrize in an additive fashion using a modulo operation due to the fact that hue is represented as a rotation angle in HSV space. On the other hand, saturation and value are discontinuous parameters that vary between 0 and 1, and cannot be directly modeled in an additive fashion, as they can't take values outside their range. Instead, we model them as multiplicative factors in R + , where we first exponentiate η s and η v to ensure the multiplicative factors are 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.In order to ensure that we can easily backpropagate through the clipping operation, we define a passthrough_clip function in Jax, where we define a custom gradient that doesn't zero out gradients even if the inputs to the function are out of bounds. We find that using the passthrough_clip operation is essential to training the model.The experiments for this paper were performed on a cluster equipped with NVIDIA A100 GPUs. All model training requires only a single such GPU. However, we used up to 64 GPUs at a time to run our hyper-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 to 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-parameter sweeps for some figures are reused for others, as discussed in Appendix C. In this section, we compare the ability of our method to learn symmetries to LieGAN [Yang et al., 2023] , which uses a generator-discriminator framework to automatically discover equivariances from a dataset using generative adversarial training. Similar to [Yang et al., 2023] , we transform the MNIST dataset to have rotations in the range [-45 • , 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 to c = 1, and consider learnable 6-dimensional Lie matrices in the generator model. The discriminator model consists of a pre-trained LeNet5 feature extractor as the backbone, and the validator is a 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 L = 0.02 -0.34 0.28 0.33 0.08 -0.05 0 0 0 . In Figure 15 , we can see that LieGAN struggles to correctly recover the range of invariances present in the training dataset, especially for translations in x. It is also unable to provide a fine-grained representation of 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 a 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 be 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 in the dataset and those learned by LieGAN. For example, in their Figure 11 , we see that the sampled digits are often rotated by significantly more than 45 • . Furthermore, we see evidence of typical GAN mode collapse, with many very similar rotations for each digit.In this section, we provide a "negative" result for our SGM when applied to the PatchCamelyon dataset [Veeling et al., 2018] . The examples in this dataset, unlike those used in Section 4, contain "content" up to the boundaries of the images. Figure 16 shows examples of the prototypes and learned distributions for this dataset, with affine and color transformations. In particular, the allowed rotation was between ±180 • , while the actual dataset has only a rotational invariance of ±n × 90 • . We see that in some cases the prototypes are rotated by close to ±n × 45 • relative to the original images. In other cases, the rotation of the prototypes relative to the original images is closer to ±n × 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 is 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.(a) Top: samples from the test set. Mid: prototypes for each test example. Bot: resampled versions of each test example given the prototype.- Figure 17 compares a standard VAE with AugVAE, an SGM-VAE hybrid model. We see that for small amounts of data, the VAE and AugVAE 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.EQUATIONThis "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.In this section, we provide additional plots to supplement those in Section 4.Figure 18 extends the results in Figure 11 to by including an 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 is able to provide the same or very similar prototype for observations in the same orbit. Secondly, it provides many more resampled examples of each digit, further demonstrating that our SGM has correctly captured the symmetries present in the dataset. Figure 20 expands on Figure 8c in the same way.Figure 21 extends Figure 9 by including all of the digits shown in Figure 19 . The conclusions are much the same as 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 is indeed the case. We did not control the value of the images, so it is more difficult to interpret those. However, given that most (non-black) pixels are bright (i.e., close to 1) it makes sense that our SGM learns multiplicative values closer to 1.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. - 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, we see that the empirical distributions match in terms of ranges, although they are imperfect in density.1 2 0 1 2 -1 2 0 1 2 -π 0 π -1 0 1 -1 0 1• While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a 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 important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.If fω is equivariant by construction, our SSL scheme is unnecessary. Alas, such constructions are unknown for many transformations, like those in this paper. Thus, we provide a general method for learning equivariances.This can occur in practice since our SSL objective-which aims to make prototypes as similar as possiblecan trivially be minimized by removing all of the contents of an image.5 Because '8' is symmetric, p (η | x, x) could be any convex combination of the two delta distributions. However, for a more realistic example, consider a prototype '8' with a smaller upper loop. In this case, the p (η | x) must be bimodal to capture '8's with both smaller upper and lower loops.We use rotation ∼ U (-15 • , 15 • ), zoom ∼ U (-10%, 10%), and x/y-shift ∼ U (-2px, 2px).Our GalaxyMNIST results have the same issue as our dSprites results-the sweep included a consistency loss multiplier which was always set to a value of 1 in our experiments. This results in some small performance degradations.Our PatchCamelyon results have the same consistency multiplier issue as our dSprites and GalaxyMNIST results. 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: Intuition behind the objective.
Experimentally 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.

Types of transformations.
The 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?
questions: 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?
2. 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?
3. 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: limitations: The paper offers a strong contribution proposing a probabilistic generative model that describes that as coming from transformed latent prototypes.
paper_summary: paper_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.
strength_weakness: strength_weakness: Some parts of the paper are confusing to me. 

What 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. 

I 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.

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

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.  

I 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.
questions: questions: Please address the weaknesses I have listed above.
limitations: limitations: I don't think the paper mentions the limitations explicitly, which the authors should try to.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

2. 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.

3. 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.

4. 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.

5. 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.

6. 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.
questions: questions: Please see the above section
limitations: limitations: Yes
paper_summary: paper_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.
strength_weakness: strength_weakness: * 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?
* 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.
* (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.
* 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].
questions: 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?

## References

[1] Benton, Gregory, et al. "Learning invariances in neural networks from training data." Advances in neural information processing systems 33 (2020): 17605-17616.

[2] Falorsi, Luca, et al. "Reparameterizing distributions on lie groups." The 22nd International Conference on Artificial Intelligence and Statistics. PMLR, 2019.
limitations: limitations: The authors have discussed the limitations.

Comments: title: End of discussion period summary
Thank you for the response
Brief discussion of additional related work for reviewer TzHS
comment: Thank you to all of the reviewers for engaging with our rebuttal and the discussion period. We are very happy that two of the reviewers (TzHS and d33o) have raised their scores, that the other two reviewers (rf86 and GGon) maintain their recommendations to accept the paper, and that *the scores now lean towards an accept on average*. 

We'd like to highlight a few of the comments that came up during the discussion period:
* "I regard this a **technically strong paper, with novel ideas and a good execution**, and therefore keep my recommendation for acceptance with score rating 8." – Reviewer rf86
* "... **my main concern about demonstrating the correct conditional distribution learning [...] has been addressed**, I will substantially raise the score..." – Reviewer d33o
* "**Most of my concerns have been addressed**... Overall, I maintain my opinion that this is a **well-written paper with clear motivation and reasonable experiment results**." – Reviewer GGon

We believe that these comments highlight that the core concerns of the reviewers have been addressed, and that the reviewers see our paper as publication worthy due to it's novelty, technical strength, and strong presentation. 

The main remaining request is to include experiments on another (more complicated) dataset. While we haven't been able to do so due to the tight timelines and compute limitations in this discussion period, we will aim to include experiments based on the suggestion of reviewer TzHS in the camera ready version.
Thank you for your engagement with the review process and for increasing your score as a result – we are very grateful!

We'll make sure to incorporate your feedback. Specifically, to
* include some more of the architecture and training details in the final manuscript,
* try our best to run some additional experiments with the datasets you have suggested (PatchCamelyon looks particularly interesting),
* include your suggested reference (Mumford and Desolneux), and
* include a dedicated paragraph on the limitations we have discussed in our previous response.
Thank you for your strong endorsement of our paper. We will be sure to include all of the answers to your questions as additional discussion on the final camera-ready manuscript.
The responses are indeed helpful in understanding some parts of the paper that I didn't before.

I understand that the architecture details and other training details are in the appendix. Yes, it would be helpful if some of the important details are moved to the main paper.

I understand the connection to existing work better.

I still think having one dataset with more challenging images is important. At the very least, it shows what the limitations of the method are. Even a dataset like CIFAR-100 with 32 x 32 color images or CUB-200-2011 or PatchCamelyon would be interesting to show some experimental results on.

A separate paragraph on limitations would be good to have rather than a footnote. For example, the authors say that running their method on more challenging datasets is computationally difficult. I don't understand if this is because the method is inherently more complex because of self-supervised learning objectives etc. or some computational constraints the authors may face.

This is not too important, but this is one reference for Pattern Theory: Pattern Theory, the Stochastic Analysis of Real World Signals by Mumford and Desolneux. But the authors should not feel compelled to include this is in the paper if they don't believe it to be very related. 

Overall, I think the reviewers have addressed my concerns to some extent. I think it is important to have some experimental validation of more challenging datasets. Based on the above, I will raise my score to a weak accept.
We thank the authors for the further clarifications and answering my answers and hope these are included as discussions the final manuscript. I regard this a technically strong paper, with novel ideas and a good execution, and therefore keep my recommendation for acceptance with score rating 8.
> While the authors provide a general method for learning equivariant, it would be beneficial to compare it with a hardwired equivariant to show its competitive performance.

We want to stress that we make no claims that our method for learning an equivariant f_w will match the performance of a hardwired f_w, in fact, we acknowledge in the paper that having some equivariances hardwired would likely lead to an increase in performance.

However, we agree that comparison with a hardwired equivariant $f_\omega$ would be interesting. If you have suggestions or know of relevant work for how to parameterise such an $f_\omega$, we'd be happy to run those experiments. This would have been easy if $f_\omega$ was a function from image to image space equivariant to translations/rotations, but we don't know of any work that does this for functions from an image space to transformation-parameter space.

> I remain unconvinced that boundary effects are insignificant. The method’s core idea, when used as a preprocessing step for VAE, is to augment samples in the group orbit accurately. Although GalaxyMNIST does not vanish at the boundary, it primarily consists of sparse objects on a black background.

As we mentioned in our previous response, we acknowledge that this is a potential limitation of our method, but note that this is also a limitation of existing published methods. Thus, we hope that we will not be held to a higher standard for publication.

> The paper assumes prior knowledge of the group symmetry in the dataset. If I understand correctly, the comparison to [Yang et al. 2023] might be unfair, as [Yang et al. 2023] genuinely learns the general linear group, whereas SGM assumes the transformation group is limited to the rotation and scaling subgroup.

To clarify, in our experiments for this paper our SGM covers rotation, scaling, *and shifting*. This makes it slightly less flexible than the LieGAN of Yang et al. [2023], since their method is also able to learn shearing and flipping. However, we feel that the comparison is largely fair, since we are not reporting any quantitative results, and instead focus on qualitative comparisons. From these qualitative comparisons, it is clear that fundamental differences between our two approaches (e.g., our SGM learning conditional distributions) are the dominant reason for different behavior, rather than the choice of assumed transformation group. We also note that our SGM is also capable of learning shearing and flipping, we simply didn't include these transformations in our experimental results. However, expending our assumed transformation group to include these is trivial. 

> Nonetheless, since my main concern about demonstrating the correct conditional distribution learning (points 1, 2, and 6 in my original review) has been addressed, I will substantially raise the score while still reserving my opinion on other aspects.

Thank you for increasing your score and engaging with the rebuttal. We very much appreciate the discussion and we are glad that the rebuttal has already addressed many of your concerns. We hope that we have addressed your remaining concerns and that you will consider increasing it further.
> Regarding your second point, I agree that the proposed method can improve sample efficiency and model fit by decoupling a small number of variations described by symmetry. I just feel that since the orbital distributions are low-dimensional, using a complex generative model may not be the most simple and efficient way to achieve this. This is a somewhat subjective matter, and I will not make further comments.

While these distributions are low-dimensional, they can often be complex, with dependencies between the different dimensions. For instance, in the case of affine transformations, there is a non-trivial relationship between rotation and shifting/scaling. As a result, we found that using a coupled neural spline flow was actually required for accurately modeling the distribution over transformation parameters. In cases where the dimensions do not interact it is possible to greatly simplify the model (e.g., a flow with a single layer and no coupling). In other words, we view the complexity of the models for transformation parameter distributions as a problem-specific design choice rather than an inherent feature of our method.
We thank the reviewer again for the effort they put into reviewing our paper. Since there are only a few working days left for the discussion period, we would like to ask if our response satisfied the reviewer's concerns. If that is the case, we kindly invite them to raise their score. If there are still any remaining concerns, we are happy to discuss them here.
We thank the reviewer again for the effort they put into reviewing our paper. Since there are only a few working days left for the discussion period, we would like to ask if our response satisfied the reviewer's concerns. If that is the case, we kindly invite them to raise their score. If there are still any remaining concerns, we are happy to discuss them here.
I thank the authors for their detailed response. After reviewing the paper and the feedback, I acknowledge that my previous review had flaws. Specifically, Figures 14, 19, and 20 in the appendix convincingly demonstrate that the proposed method learns the correct conditional distribution, and these figures slipped through my first review. Since my primary concern was the lack of clear evidence that the method achieves this rather than merely **augmenting** training data and **displaying** visually realistic samples, I find it necessary to raise the score.

However, I still have several concerns:

1. While the authors provide a general method for learning equivariant $f_w$, it would be beneficial to compare it with a hardwired equivariant $f_w$ to show its competitive performance.

2. I remain unconvinced that boundary effects are insignificant. The method’s core idea, when used as a preprocessing step for VAE, is to augment samples in the group orbit accurately. Although GalaxyMNIST does not vanish at the boundary, it primarily consists of sparse objects on a black background.

3. The paper assumes prior knowledge of the group symmetry in the dataset. If I understand correctly, the comparison to [Yang et al. 2023] might be unfair, as [Yang et al. 2023] genuinely learns the general linear group, whereas SGM assumes the transformation group is limited to the rotation and scaling subgroup.

Nonetheless, since my main concern about demonstrating the correct conditional distribution learning (points 1, 2, and 6 in my original review) has been addressed, I will substantially raise the score while still reserving my opinion on other aspects.
Thank you for your detailed response. Most of my concerns have been addressed.

Regarding your second point, I agree that the proposed method can improve sample efficiency and model fit by decoupling a small number of variations described by symmetry. I just feel that since the orbital distributions are low-dimensional, using a complex generative model may not be the most simple and efficient way to achieve this. This is a somewhat subjective matter, and I will not make further comments.

Overall, I maintain my opinion that this is a well-written paper with clear motivation and reasonable experiment results. I will keep my current recommendation.
*Group Equivariant Subsampling*. This paper introduces group-equivariant sub/upsampling layers. These layers can be used to construct group-equivariant CNNs with pooling-like layers. Our paper does not develop any novel architectures for group-equivariance, though our inference network could be made equivariant, as discussed in “Invariance of fω and the prototypes”. Thus, this line of work is related but orthogonal to ours.

*Learning Equivariances and Partial Equivariances from Data*. Like ours, this work focuses on learning partial symmetries from data. However, this work is more closely related to that of  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], Immer et al. [2022, 2023], Miao et al. [2023], and Mlodozeniec et al. [2023], in that it is concerned with the supervised learning regime.

*Deforming autoencoders: Unsupervised disentangling of shape and appearance*. Like ours, this paper focuses on unsupervised learning of prototypical objects (called ‘templates’ in their work) and transformations to convert such prototypes into observations. However, our work is more general, being applicable to a more general class of transformations. Conversely, theirs includes many interesting inductive biases specific to shape transformations. Their method is not probabilistic, meaning that they are unable to sample novel observations. 

*Disentangling images with Lie group transformations and sparse coding*. Like ours, this paper learns to disentangle images into prototypes and their transformations. However, theirs differs from ours in two key ways. Firstly, they focus only on transformations corresponding to Lie groups whereas ours is more general. Secondly, they learn prototypes that are themselves composed of sparse combinations of learned dictionary elements. We see this work as complementary to ours, as the core idea of combining sparse coding with learned transformations could also be applied to our method.