A Generative Model of Symmetry Transformations

For ${\mathcal{T}}_{\bm{\upeta}}^{\scalebox{0.5}[1.0]{$\>-$}1}$, with ${\bm{\upeta}}$ given by $f_{\bm{\upomega}}$, to map ${\mathbf{x}}$ to a prototype $\hat{\mathbf{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 ${\mathcal{T}}^{\scalebox{0.5}[1.0]{$\>-$}1}_{f_{\bm{\omega}}({\bm{x}})}({\bm{x}})$ should be invariant to transformations ${\mathcal{T}}_{{\bm{\eta}}^{\prime}}$ of ${\bm{x}}$:

To learn such a function, we optimize for this property directly. To this end, we sample transformation parameters ${\bm{\eta}}_{\text{rnd}}$ from some distribution over parameters $p({\bm{\upeta}}_{\text{rnd}})$. This allows us to get random samples ${\bm{x}}_{\text{rnd}}:={\mathcal{T}}_{{\bm{\eta}}_{\text{rnd}}}({\bm{x}})\in{\mathcal{X}}$ in the orbit of any given element ${\bm{x}}\in{\mathcal{X}}$. Since we want full (i.e., strict) invariance, $p_{\>\!\!}\left({\bm{\upeta}}_{\text{rnd}}\right)$ must have support on the entire orbit (van der Ouderaa and van der Wilk, 2022). We then learn an equivariant function $f_{\bm{\upomega}}$³³3 If $f_{\bm{\upomega}}$ 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. via a self-supervised learning (SSL) scheme 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:

Our actual objective differs slightly. Since ${\mathcal{T}}_{{\bm{\eta}}^{\prime}}({\bm{x}}^{\prime})={\mathcal{T}}_{{\bm{\eta}}^{\prime\prime}}({\bm{x}}^{\prime\prime})$ implies ${\bm{x}}^{\prime}={\mathcal{T}}^{\scalebox{0.5}[1.0]{$\>-$}1}_{{\bm{\eta}}^{\prime}}\circ{\mathcal{T}}_{{\bm{\eta}}^{\prime\prime}}({\bm{x}}^{\prime\prime})$, we use

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 Algorithm˜1 of Algorithm˜1.

Once we have a prototype inference function, we simply learn $p_{{\bm{\psi}}\!}\left({\bm{\upeta}}\,\middle|\,\hat{\mathbf{x}}\right)$ by maximum likelihood on the created data pairs $\left\{f_{\bm{\upomega}}({\bm{x}}_{i}),{\mathcal{T}}^{\scalebox{0.5}[1.0]{$\>-$}1}_{f_{\bm{\upomega}}({\bm{x}}_{i})}({\bm{x}}_{i})\right\}$. This is shown in ˜6 of Algorithm˜1. While we need to specify the kinds of symmetry transformations ${\mathcal{T}}_{\bm{\upeta}}$ we expect to see in the data, by learning $p_{{\bm{\uppsi}}\!}\left({\bm{\upeta}}\,\middle|\,\hat{\mathbf{x}}\right)$ 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_{\>\!\!}\left({\bm{\upeta}}_{\text{rnd}}\right)$) 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 ${\bm{x}}$, and (b) resample new “naturally” augmented versions of the element, we do not need to learn $p_{\>\!\!}\left(\hat{\mathbf{x}}\right)$. We can do (a) by querying $p_{\>\!\!}\left({\bm{\upeta}}\,\middle|\,\hat{\mathbf{x}}=\hat{\bm{x}}\right)$ for $\hat{\bm{x}}:={\mathcal{T}}^{\scalebox{0.5}[1.0]{$\>-$}1}_{f_{\bm{\eta}}({\bm{x}})}({\bm{x}})$, and we can do (b) by sampling ${\bm{\eta}}\sim p_{\>\!\!}\left({\bm{\upeta}}\,\middle|\,\hat{\bm{x}}\right)$ and transforming $\hat{\bm{x}}$ to get ${\bm{x}}:={\mathcal{T}}_{\bm{\eta}}\left(\hat{\bm{x}}\right)$. Of course, if one wanted to sample new prototypes, one could fit $p_{{\bm{\uptheta}}\!}\left(\hat{\mathbf{x}}\right)$ using, e.g., a VAE. Not learning $p_{\>\!\!}\left(\hat{\mathbf{x}}\right)$ greatly simplifies training for complicated datasets that would otherwise require a large generative model, an observation made by Dubois et al. (2021).

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({\bm{\upeta}})$ where composition can be performed—e.g., as matrix multiplication ${\mathcal{T}}_{{\bm{\eta}}_{2}}\circ{\mathcal{T}}_{{\bm{\eta}}_{1}}={\mathcal{T}}^{\prime}_{T({\bm{\eta}}_{2})T({\bm{\eta}}_{1})}$, in the case where $T$ is a group representation—then we recommend composing transformations in that space to minimize the number of applications.

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 $\hat{\mathbf{x}}$ exists for any ${\mathbf{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-bounds⁴⁴4 This can occur in practice since our SSL objective—which aims to make prototypes as similar as possible—can trivially be minimized by removing all of the contents of an image. ). If appropriate bounds are known a priori, we can prevent severe information loss by constraining ${\bm{\upeta}}_{\text{min}}$ and ${\bm{\upeta}}_{\text{max}}$ using $\mathtt{tanh}$, $\mathtt{scale}$, and $\mathtt{shift}$ bijectors. Alternatively, we can augment the SSL loss in Algorithm˜1 with an invertibility loss

In practice, our transformation inference network $f_{{\bm{\upomega}}\!}\left({\mathbf{x}}\right)$ will not be perfect; see Figure˜10. Even after training, there may be small variations in the prototypes $\hat{\mathbf{x}}$ corresponding to different elements in the orbit of ${\mathbf{x}}$. To make $p_{{\bm{\psi}}\!}\left({\bm{\eta}}_{\mathbf{x}}\,\middle|\,\hat{\mathbf{x}}\right)$ 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_{{\bm{\psi}}\!}\left({\bm{\eta}}_{\bm{x}}\,\middle|\,{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}\pgfsys@color@cmyk@stroke{1}{1}{0}{0}\pgfsys@color@cmyk@fill{1}{1}{0}{0}\hat{\bm{x}}^{\prime}}\right)$, where $\hat{\bm{x}}^{\prime}={\mathcal{T}}^{\scalebox{0.5}[1.0]{$\>-$}1}_{f_{\bm{\upomega}}({\mathcal{T}}_{{\bm{\eta}}_{\text{rnd}}}({\bm{x}}))}({\mathcal{T}}_{{\bm{\eta}}_{\text{rnd}}}({\bm{x}}))$ as in our SSL objective. Averaging the loss over multiple samples—e.g., 5—of ${\bm{\upeta}}_{\text{rnd}}$ is beneficial.

Consider a dataset of ‘8’s rotated in the range $-30^{\circ}$ to $30^{\circ}$: { 8 , …, 8 , …, 8 }. Let us assume that the prototype is ‘8’. Figure˜5(a) shows $p_{\>\!\!}\left({\upeta}\,\middle|\,{\mathbf{x}},\,\hat{\mathbf{x}}\right)$, an example of the true distribution for ${\upeta}$ given ${\mathbf{x}}$ and $\hat{\mathbf{x}}$, for several observations, under the data generating process⁵⁵5 Because ‘8’ is symmetric, $p_{\>\!\!}\left({\upeta}\,\middle|\,{\mathbf{x}},\,\hat{\mathbf{x}}\right)$ 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_{\>\!\!}\left({\bm{\upeta}}\,\middle|\,\hat{\mathbf{x}}\right)$ must be bimodal to capture ‘8’s with both smaller upper and lower loops. . These distributions are composed of deltas because only certain values of ${\upeta}$ will transform $\hat{\mathbf{x}}$ into ${\mathbf{x}}$. Figures˜5(b) and 5(c) compare idealised examples of the learned $p_{{\bm{\uppsi}}\!}\left({\upeta}\,\middle|\,\hat{\mathbf{x}}\right)$—given a simple uni-modal Gaussian family and a more flexible bi-modal mixture-of-Gaussian family—with the aggregate true distribution $p_{\>\!\!}\left({\upeta}\,\middle|\,\hat{\mathbf{x}}\right)=\sum_{{\mathbf{x}}\in\{{\color[rgb]{0.15,1,0.67}\definecolor[named]{pgfstrokecolor}{rgb}{0.15,1,0.67}\pgfsys@color@cmyk@stroke{0.85}{0}{0.33}{0}\pgfsys@color@cmyk@fill{0.85}{0}{0.33}{0}\rotatebox[origin={c}]{30.0}{8}},\ldots,{\color[rgb]{0.0600000000000001,0.46,1}\definecolor[named]{pgfstrokecolor}{rgb}{0.0600000000000001,0.46,1}\pgfsys@color@cmyk@stroke{0.94}{0.54}{0}{0}\pgfsys@color@cmyk@fill{0.94}{0.54}{0}{0}\rotatebox[origin={c}]{0.0}{8}},\ldots,{\color[rgb]{0.21,0.12,1}\definecolor[named]{pgfstrokecolor}{rgb}{0.21,0.12,1}\pgfsys@color@cmyk@stroke{0.79}{0.88}{0}{0}\pgfsys@color@cmyk@fill{0.79}{0.88}{0}{0}\rotatebox[origin={c}]{-30.0}{8}}\}}p_{\>\!\!}\left({\upeta}\,\middle|\,{\mathbf{x}},\,\hat{\mathbf{x}}\right)$. 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 by modeling $p_{{\bm{\uppsi}}\!}\left({\upeta}\,\middle|\,\hat{\mathbf{x}}\right)$ with normalizing flows, which can match a wide range of distributions.

Consider a dataset of ‘2’s and ‘8’s rotated in the range $-30^{\circ}$ to $30^{\circ}$: { 2 , …, 2 , …, 2 , 8 , …, 8 , …, 8 }, with prototypes ‘2’ and ‘8’. Figure˜6(a) shows $p_{\>\!\!}\left({\upeta}\,\middle|\,{\mathbf{x}},\,\hat{\mathbf{x}}\right)$, an example of a true distribution over ${\upeta}$, for several observations. Figures˜6(b) and 6(c) compare idealized examples of learned distributions over ${\upeta}$ and ${\upeta}\,|\,\hat{\mathbf{x}}$. Without dependence on $\hat{\mathbf{x}}$, the model must place probability mass between $-150^{\circ}$ and $150^{\circ}$, in order to capture the symmetries of the ‘8’s, however this results invalid digits—such as { 2 , 2 , 2 }—which do not come from true data distribution. On the other hand, when ${\upeta}$ depends on $\hat{\mathbf{x}}$, the distribution conditioned on the prototype for the ‘2’s only needs to place mass in $[-30^{\circ},30^{\circ}]$.

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^{\circ},\,180^{\circ}]$ 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 $\hat{\mathbf{x}}$ contains no information about the symmetries in the data.

Consider a dataset of ‘2’s rotated in the range $-30^{\circ}$ to $30^{\circ}$: { 2 , …, 2 , …, 2 }. Figure˜7(a) shows predicted prototypes and the corresponding distributions over ${\upeta}$ for several observations. There are three cases: (a) a fully-invariant $\hat{\mathbf{x}}$, i.e., there is a single prototype, (b) a partially-invariant $\hat{\mathbf{x}}$, for which there are two prototypes in this example, and (c) a non-invariant $\hat{\mathbf{x}}$, which takes the partially-invariant 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_{{\bm{\uppsi}}\!}\left({\bm{\upeta}}\,\middle|\,\hat{\mathbf{x}}\right)$ does not represent the distribution of naturally occurring transformations of $\hat{\mathbf{x}}$ in the data. This is illustrated in Figures˜7(b), 7(c) and 7(d), which show idealized examples of the learned $p_{{\bm{\uppsi}}\!}\left({\upeta}\,\middle|\,\hat{\mathbf{x}}\right)$ in each case. While the distribution in Figure˜7(b) matches the distribution of transformations in the dataset, in Figures˜7(c) and 7(d) 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 $\scalebox{0.5}[1.0]{$\>-$}90^{\circ}$. Assuming we can find a prototype $\hat{\bm{x}}$, we would like $p_{\>\!\!}\left({\upeta}\,\middle|\,\hat{\mathbf{x}}=\hat{\bm{x}}\right)$ to represent all naturally occurring augmentations. Unless $\hat{\mathbf{x}}$ is unique, this won’t necessarily be the case, as shown in Figures˜7(c) and 7(d).

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 several illustrative examples which warrant further discussion. The heart sprites in Figure˜8(a) 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˜8(b) 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’.

In ${\bm{\eta}}$, each dimension corresponds to a different transformation. We refer to $p_{{\bm{\uppsi}}\!}\left({\bm{\upeta}}_{i}\,\middle|\,{\mathbf{x}}\right)$ as the marginal distribution of a single transformation parameter. Figure˜9 shows the learnt marginal distributions for several digits from Figure˜8(b). We see that each of the parameters has its own range and shape. 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^{\text{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_{{\bm{\uppsi}}\!}\left({\bm{\eta}}\,\middle|\,{\mathbf{x}}\right)$, as conjectured in Section˜3.2. We also find significant dependencies between dimensions of ${\bm{\upeta}}$ (e.g., rotation and translation in dSprites).

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 $\hat{\mathbf{x}}_{1}$ is most different from the previous $\hat{\mathbf{x}}_{0}={\mathbf{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.