Extended Abstract Track
Learning Deep Generative Models with Invariance under Symmetry Transformations
Editors: Sophia Sanborn, Christian Shewmake, Simone Azeglio Arianna Di Bernardo Miolane
Abstract
While imbuing a model with invariance under symmetry transformations can improve data efficiency and predictive performance, most methods require specialised architectures and, thus, prior knowledge of the symmetries.
Unfortunately, we don’t always know what symmetries are present in the data.
Recent work has solved this problem by jointly learning the invariance (or the degree of invariance) with the model from the data alone.
But, this work has focused on discriminative models.
We describe a method for learning invariant generative models.
We demonstrate that our method can learn a generative model of handwritten digits that is invariant to rotation.
1. Introduction
Generative models are often motivated as being more data efficient than discriminative models (Welling, 2019).
We expect that since we are encoding our beliefs about how the data are generated (our inductive biases) into the model, it will be able to learn from fewer data.
Unfortunately, deep generative models are more data hungry than their discriminative counterparts.
This is because deep generative models tend to only incorporate very general inductive biases (smoothness, hierarchies of random variables, etc.).
One kind of inductive bias that does improve efficiency is invariance under symmetry transformations in the data (van der Ouderaa and van der Wilk, 2022).
In some cases, we know apriori which transformations the model should be invariant to.
However, in many cases, we need to learn when (or to what extent) to be invariant.
For example, in the case of handwritten digit recognition, we know that we should be invariant to some amount of rotation.
However, a model that is invariant to rotations in the full range [ − 180 ◦ , 180 ◦ ] would be unable to distinguish between ‘6’ and ‘9’.
Thus, the degree of invariance should be learned.
Recent work has focused on solving this problem for discriminative models (Nalisnick and Smyth, 2018; van der Wilk et al., 2018; Benton et al., 2020; Schw¨obel et al., 2021; van der Ouderaa and van der Wilk, 2022; Immer et al., 2022).
But, the problem has been neglected for deep generative models – our focus in this paper.
We aim to construct a deep generative model p ( x ) which is invariant to symmetry transformations of x . The generative model Extended Abstract Track
should be equipped with invariant representations ˆ x containing no information about the transformation.
Concretely, p (ˆ x | x ) = p (ˆ x | T η ( x )), where η is a random variable which parameterizes the transformation T .
For example, T could be the rotation of an image, and η could be the angle.
The degree of invariance learnt—i.e., η —should match the true invariance of the data.
That is, unlike Benton et al.
(2020) we should not learn a ‘maximum data augmentation’.
2. The Model and Design Choices
[Equation]
N
[Figure]
We assume the following generative model for the data x , also depicted in Figure 1:
[Equation]
where the prototype ˆ x can be considered a standardised or reference example with no transformation applied to it.
For example, this could be a ‘6’ at in the upright (90 ◦ ) rotation 1 .
Here, p θ (ˆ x ) is some generative model with parameters θ — a Variational Autoencoder (VAE) or a normalizing flow, for example.
Similarly, p ψ ( η | ˆ x ) is a neural network (NN), with parameters ψ , that predicts the distribution over transformations for a given prototype. We choose p ( x | ˆ x , η ) = N ( T η (ˆ x ) , σ I ).
Note that while we do need to specify the kinds of symmetry transformations T we expect to see in the data, by learning η the model will control the degree to which it is invariant. Thus, we can specify several potential symmetry transformations, and learn not to be invariant to those not actually present in the data. E.g., we could specify T η as a family of affine transformation, as in Benton et al. [Equation]
where G i are infinitesimal generator matrices for various affine transformations. For all transformations without symmetries in the data [Equation] we expect to learn p ( η i | ˆ x )
≈ δ
( η i − 0)
[Equation]
We propose a VAE-like inference strategy for this model.
Given an inference model— depicted in Figure 1—we can derive an Evidence Lower BOund (ELBO) for jointly learning the generative { θ , ψ } and inference { ϕ , ν } parameters
[Equation]
with gradient descent:
[Equation]
Extended Abstract Track
where q ν ( η | x ) and q ϕ (ˆ x | x ) are NNs with parameters ν and ϕ .
By choosing q ϕ (ˆ x | x ) to be a function invariant under T η , we encourage ˆ x not to contain any information about η .
Perhaps counter-intuitively, and in contrast with previous approaches to learning invariances (van der Wilk et al.,
; Benton et al., 2020; Schw¨obel et al., 2021; van der Ouderaa and van der Wilk, 2022; Immer et al., 2022), we do not make this NN partially invariant.
Nonetheless, as we will show, we can still learn a degree of invariance.
We now present three informal conjectures to explain some of the design choices for the model in Figure 1.
We illustrate these conjectures by means of examples.
In each example we will assume that T η is counter-clockwise rotation, and thus η is the angle of rotation.
Conjecture 1: p ψ ( η | ˆ x ) must be sufficiently flexible.
Consider a dataset of slightly rotated ‘8’s: { 8, 8 , 8 } .
Let us assume that the prototype is ‘8’.
Table 1 shows p ( η | x , ˆ x ), the true distribution for η given x and ˆ x , under the data generating process.
Because ‘8’ is symmetric, p ( η | x , ˆ x ) is the sum of two deltas.
Figure 2 compares the learned p ψ ( η | ˆ x ) given a simple uni-modal Gaussian family and a more flexible bi-modal mixture-of-Gaussian family with the aggregate true distribution p ( η | ˆ x ) = (cid:80) x ∈{ 8 , 8 , 8 } p ( η | x , ˆ x ) 2 .
Here, the simple distribution is not flexible enough, which results in a large amount of probability mass being wasted on angles with low density under the true data generating process.
The definition of ‘flexible enough’ will, of course, vary depending on the problem. [Table]
[Figure]
Conjecture 2: q ϕ (ˆ x | x ) must be fully invariant w.r.t η .
Consider a dataset of slightly rotated ‘2’s: { 2, 2 , 2 } .
Assume that the prototype is predicted by a partially invariant q ϕ (ˆ x | x ). That is, q ϕ (ˆ x | x ) is only similar for ‘2’s which are equivalent under rotations smaller than ± ρ .
Table 2 shows the predicted prototypes and the corresponding distributions over η for two cases: ρ ≥ 60 ◦ and 60 ◦ > ρ .
The threshold is 60 ◦ because each ‘2’ can be transformed into each other ‘2’ by a rotation of ± 60 ◦ or less.
In the first case, each ‘2’ maps to the same prototype.
However, in the second case, each ‘2’ can have a unique prototype.
Thus, the prototype contains rotation information, which is invalid under our generative model of the data.
Unfortunately, as illustrated by Figure 3, in the second case p ψ ( η | ˆ x ) can become arbitrarily large, which corresponds to a higher marginal likelihood [Equation]
Extended Abstract Track
assuming that p ( x | η, ˆ x ) and p θ (ˆ x ) are roughly the same in both cases 3 .
By making q ϕ (ˆ x | x ) fully invariant, the prototypes will be the same, and we avoid this degenerate optima.
[Table]
[Figure] [Table] [Figure] [Table]
[Figure]
Conjecture 3: The distribution over η must depend on ˆ x.
Consider a dataset of slightly rotated ‘2’s and ‘8’s: { 2, 2 , 2 , 8, 8 , 8 } , with prototypes ‘2’ and ‘8’.
Table 3 shows p ( η | x , ˆ x ), the true distribution over η .
Figure 4 compares learned distributions over η and η given ˆ x .
Note that in the first case, we are able to sample invalid digits such as { 2, 2 , 2 } , which did not appear in the dataset.
However, when η depends on ˆ x , this does not occur. [Table]
[Figure]
3. Experiments
We validate that our model is able to learn invariance under symmetry transformations by training on a rotated version of the MNIST dataset.
We choose p θ (ˆ x ) to be a VAE, and we implement p ψ ( η | ˆ x ) as a Neural Spline Flow (Durkan et al., 2019) f ψ (ˆ x ) with base distribution N ( µ ψ (ˆ x ) , σ ψ (ˆ x )) to ensure sufficient flexibility. q ν ( η | x ) is defined similarly. We do not place any restrictions on η , allowing it be learnt freely 4 .
Further implementation details are in Appendix A. Figure 5 shows that our model successfully (i) learns prototype images that are invariant to rotations—each ˆ x is at the same arbitrary angle
≈ 135 ◦ regardless of the orientation of x —and (ii) reconstructs rotated digits in the correct orientation.
Figure 6 shows that the model behaves sensibly in the presence of different transformation intensities—as the digits are rotated more, the model predicts a wider range of angles.
Note that as the degree of rotation in the data increases, so does the range of angles predicted by our model.
This is in contrast to methods such as Benton et al.
(2020) which learn a ‘maximum data augmentation’ which decreases as the data is transformed to a greater extent.
4. Conclusion
[Figure]
We have presented a method for learning a generative model with representations that are invariant to symmetry transformations.
Our method learns these invariances in an unsupervised manner, without signal from an auxiliary task such as classification.
We discussed various pitfalls for invariant generative models, which motivated our design choices.
We also provided initial qualitative results, showing that the model can learn invariant representations for rotated MNIST digits and behaves reasonably under varying amounts of rotation.
However, this is all a first step towards learning invariant generative models.
Important next steps include (i) formalization of our motivating conjectures, (ii) quantitative results with comparisons to baseline methods, (iii) extension to a wider range of transformations (e.g., other affine transformations, colour transformations), (iv) ablation studies for all of our design choices, (v) scaling up to more interesting problems than MNIST digits, and (vi) investigating our motivating example of improved data efficiency for invariant generative models.
Extended Abstract Track
Acknowledgments
JUA acknowledges funding from the EPSRC, the Michael E. Fisher Studentship in Machine Learning, and the Qualcomm Innovation Fellowship.
JA acknowledges support from
Microsoft Research, through its PhD Scholarship Programme, and from the EPSRC.
SP acknowledges support from the Harding Distinguished Postgraduate Scholars Programme
Leverage Scheme.
JMHL acknowledges support from a Turing AI Fellowship under grant
EP/V023756/
Appendix A. Additional Experimental Details
[Equation]
The full generative model for our experiments, modelling p θ (ˆ x ) with a VAE, is:
[Equation]
The full graphical model is shown in Figure 7.
We use the following ELBO for jointly learning the generative parameters { θ , ψ } and the variational parameters { ϕ , ν } : log
[Equation]
We estimate the ELBO with a single Monte-Carlo sample.
We follow Benton et al.
(2020); van der Ouderaa and van der Wilk (2022);
Immer et al.
in constructing approximately invariant NNs via expectation.
Specifically, we construct an invariant encoder q ϕ ( z | x ) from a non-invariant encoder ˆ q ϕ ( z | x ):
[Equation]
Extended Abstract Track
We approximate this expectation using 10 Monte-Carlo samples.
We use jax with flax , optax , and distrax to implement our model.
Our VAE encoder q ϕ ( z | x ) is a 3-layer CNN, with { 64, 128, 256 } channels parameterising a heteroskedastic
Gaussian distribution.
The VAE decoder p θ (ˆ x | z ) is a 3-layer CNN, with transposed convolutions of { 256, 128, 64 } channels, parameterising a homoskedastic Gaussian distribution.
The latent dimension is 20.
p ψ ( η | z ) is a Neural Spline Flow with four bins, two flow layers, and base distribution N ( µ ψ ( z ) , σ ψ ( z )) where the µ ψ and σ ψ are a two-layer MLPs with hidden sizes { 64, 32 } .
q ν ( η | x ) follows the same structure.
For both flow models, each flow layer is parameterised by a 2-layer MLP of with hidden sizes { 64, 32 } .
All neural networks use relu activation functions.
The VAE encoder and decoder use
LayerNorm .
We use the adamw optimizer with a initial learning rate of 1 × 10 − 4 and a weight decay of 1 × 10 − 4 .
We train for 15000 steps with a batch size of 256.
We linearly increase the LR to 1 × 10 − 3 over the first 500 training steps and then use a cosine schedule to reduce it back to 1 × 10 − 4 by the end of training.
We multiply the VAE KL-divergence by a factor β , which has an initial value of 10 and is decayed to 1, using a cosine schedule, by the end of training.