Achieving Combinatorial Generalisation by Modeling a Group Action

# Achieving Combinatorial Generalisation by Modeling a Group Action ## Introduction Humans have the natural ability to generalise over a wide range of situations. A straightforward example of this is understanding text: You, the reader, may not have encountered this particular combination of words before, yet you are able to infer the message conveyed by this text through semantic understanding. Being able to generalise over unseen combinations of concepts is what, arguably, enables humans to communicate effectively, produce art, and invent technologies. As such, this constitutes a fundamental part of intelligence. As deep generative models move closer to replicating some of the key skills that define human intelligence, it becomes increasingly important to evaluate them using real-life human-level tasks. To date, attaining combinatorial generalisation in conventional deep learning approaches has proved difficult, as these models fail to adapt to distribution shift [[1](##References)]. Earlier research hypothesised that disentangled representation learning enables combinatorial generalisation. However, Locatello [[2](##References)] demonstrated empirically that disentangled representations are insufficient for combinatorial generalisation. Other approaches [[3,4,5,6](##References)] learned the structural causal model contained in the data to achieve combinatorial generalisation. The downside of these approaches is the necessity of expensive prior knowledge. *Achieving Combinatorial Generalisation by Modeling a Group Action* [[7](##References)] proposes a novel framework for combinatorial generalisation in deep generative models by learning to encode the underlying data transformations, as opposed to learning the data distribution. This is driven by learning the independent operations that transform a fixed image to a large variety of other images, which allows the model to apply unseen compositions of transformations during evaluation. In this blog, we aim to: 1. Describe and explain the approach of the *Modeling A Group Action* framework. 2. Visually demonstrate how the MAGANet encodes data transformations in latent space and obtain an intuition why this leads to the desired result of combinatorial generalisation. 3. Open source our reimplementation of the framework in Python using Pytorch on [Github](https://github.com/Roodster/dsait4125-cv). To the best of our knowledge an open source implementation is not available at this moment. ## Background Let's first briefly dive into the theorectical underpinnings of the approach. ### Basics of Group theory To understand the *Modeling a Group Action* framework, we will briefly address the concepts of Group theory that are at the basis of the approach. We use the notation of [[8](##References)]. **Definition:** A group is a set $G$ together with an operation * denoted as $(G, *)$ with $$(a, b) \mapsto a * b: G \times G \rightarrow G$$ satisfying the following conditions: - Associativity: The result of applying an operator in different orders is the same. $$(a * b) * c = a * (b * c).$$ - Identity: The identity property enables an element $a \in G$ to be the exact same as before applying an operator by using a specific element $e \in G$: $$a * e = a = e * a.$$ - Inverse: For all elements $a$ in the set, there has to be an 'inverse' element $a'$ such that: $$a * a' = e = a' * a.$$ **Definition:** A group action is when a group $G$ interacts with a set $X$ via function $a: G \times X \rightarrow X$, defined as $g\cdot x = \alpha(g, x)$. The group action has to follow two rules: 1. **Identity**: if $e \in G$ is the 'identity' element, then $a(e, x) = x$ for all $x \in X$. 2. **Composition**: For any pair of group elements $g, h \in G$ and $\forall x \in X$, applying $h$ then $g$ is the same as applying $gh$ to x. That is $\alpha(g, \alpha(h, x)) = \alpha(g\cdot h, x)$. ### Generative models #### Variational Autoencoders (VAE) At their core, all generative models attempt attempt to approximate $p_\theta(x)$, the underlying probability distribution that governs the data. The basic premise is that by knowing $p_\theta(x)$, we can generate data that exibits the same features/patterns as the training dataset through sampling [[9](##References)]. For Variational Autoencoders, the formulation varies slightly. Given a latent variable $Z$ with distribution $p_\theta(z)$ and a conditional distribution $p_\theta(x|z)$ where $X$ is an observation from this conditional distribution, the challenge is to learn an approximate posterior $q_\phi(z|x)$ such that the latent variables can be effectively inferred while ensuring that the generated samples from the learned latent space remain faithful to the original data distribution [[10]](##References). To optimize this approximation efficiently, VAEs maximize the Evidence Lower Bound (ELBO): $$\mathcal{L}(\theta,\phi) = \mathbb{E}_{q_\phi(z|x)}[log(p_\theta(x|z)] - \beta_{KL}D_{KL}(q_\phi(z|x)||p_\theta(z))$$ Where the first term is known as the *Reconstruction Loss*, which encourages the model to accurately reconstruct the original data and the second as the *Kullback-Leibler (KL) Divergence*, which regularizes the latent space to follow a smooth and continuous normal distribution [[10]](##References). Finally $\beta_{KL}$ determines the weight or importance assigned to the KL loss during training [[14]](##References).  #### Normalizing Flow Normalizing Flows [[11, 12]](##References) belong to a class of generative models capable of mapping distributions, often simple ones, to more complicated distributions through applying a sequence of invertible transformations to the distribution. Take for example the continuous random variable $Z \sim N(0,1)$. We define an invertible function $f: \mathbb{R^{n} \rightarrow \mathbb{R^{n}}}$ that maps a sample $z \sim Z$ to another random variable $x \sim X$, where $X = f(Z)$ and $Z = f^{-1}(X)$. Using the change of variable rule, we get that the marginal likelihood $p_{X}(x)$ is given by $$p_{X}(x) = p_{Z}(f^{-1}(x))\left|det\left(\frac{\delta f^{-1}(x)}{\delta x}\right)\right|$$ Here, $|det(\frac{\delta f^-1}{\delta x})|$ acts as a normalising term to ensure that $p_{X}(x)$ is properly adjusted for the volume changes induces by $f$. The term flow refers to ability to chain invertible transformations to create more complicated invertible transformations: $\text{f} = \text{f}_1 \circ \text{f}_2 \circ \cdots \circ \text{f}_K$. This sequence of invertible transformations is called a normalising flow. : $$\ce{x <=>[f1][f1^{-1}] h1 <=>[f2][f2^{-1}] h2 \cdots <=>[fk][fk^{-1}] z},$$ where $h_{i}$ are intermediary variables. The concepts of a normalising flow can also be used for generative modelling [[8, 11]](##References). We train a generative model $p_{\theta}(x)$ to match the true underlying data distribution $p^{*}(x)$ of the dataset $\mathcal{X}$ as closely as possible. This is done by maximising the likelihood of the observed data, which is equivalent to minimising the negative log-likelihood of a discrete i.i.d. dataset $\mathcal{D} = \{x^{(1)}, x^{(2)}, \cdots, x^{(N)}\}$: $$\mathcal{L}(\mathcal{D}) = -\frac{1}{N} \sum_{i=1}^{N}log (p_{\theta}(x^{(i)}))$$ By the change of variable formula we can compute $p_{\theta}(x)$ exactly. Since $x = f_{\theta}(z)$, we can express the log-density of $x$ in terms of the distribution and the transformations: $$log p_{\theta}(x) = log p(z) + log \left|det\left(\frac{\delta f^{-1}(x)}{\delta x}\right)\right|$$ Since the $f^{-1}$ denotes a sequence of transformations $f = f_{1} \circ f_{2} \circ \cdots \circ f_{K},$ the log determinant can be broken down into: $$log\left|det\left(\frac{\delta f^{-1}(x)}{\delta x}\right)\right| = \sum_{i=1}^{K} log \left|det\left(\frac{\delta h_{i}}{\delta h_{i-1}}\right) \right|$$ In practice, the transformations $f_{i}$ are chosen such that the Jacobian is triangular, as the determinants of these transformations are simply the sum of its diagonal elements: $$log \left|det\left(\frac{\delta f^{-1}(x)}{\delta x}\right)\right| = \sum_{i=1}^{K} log \left|diag\left(\frac{\delta h_{i}}{\delta h_{i-1}}\right) \right|$$   ## Methodology In the following section we discuss the technical implementation and describe the approach. ### Architecture At its core, MAGANet learns to reconstruct images by applying learned transformations in latent space. This is enabled by learning the transformation between two images called $x_{pivot}$ and $x_{varying}$ , top and bottom images respectively as seen in Figure 1. The pivot image is an image that is fixed during both training and testing. During training, the model learns to transform $x_{pivot}$ into $x_{varying}$. To do so the model encodes both images using: $$E(x_{pivot}, x_{varying}) = E(x_{pivot}) - E(x_{varying}) = z_{varying} - z_{pivot} = z.$$ If we use the images below as an example, where the pivot image $x_{pivot}$ is a rotated square on the top left and the varying image $x_{varying}$ is a triangle on the bottom right, $z$ encodes the transformation that transforms the shape, position, rotation, scale from $x_{pivot}$ to $x_{varying}$. Using this latent $z$ and the pivot image $x_{pivot}$, the aim is to reconstruct the image $x_{varying}$ by passing the pivot image $x_{pivot}$ and $z$ through the decoder $D$. Successful reconstruction demonstrates that the model is capable of applying the transformation $z$ to $x_{pivot}$. <center> <img src="https://hackmd.io/_uploads/SkUbESa21e.png" width="500"> </center> Figure 1: Visualisation of MAGANet architecture. Left: the pivot and varying images. Middle: The encoder and decoder modules with the latent representation inbetween. Right: Reconstructed image. The pivot image (top left) and the varying image (bottom left) are passed through the encoder $E$ to obtain the transformation latent $z$. The latent $z$ and the pivot image are fed back into the Decoder $D$ to reconstruct the image $x_{varying}$ by applying $z$ to $x_{pivot}$. #### Encoder <center> <img src="https://hackmd.io/_uploads/HyqU5S6nJg.png" width="600"> </center> Figure 2: Detailed visualisation of the Encoder module. The non-linearities between convolutional and linear layers are omitted for visualisation purposes. The encoder module is shown twice but in practice the images are passed through a shared encoder module.  The detailed encoder architecture is shown above. First, the image is passed through multiple convolutional neural layers with ReLU activation functions in between, of which the final output is flattened and passed through a 3 layered multi-layered perceptron, again with ReLU non-linearities.  In this architecture the mean $\mu$ and log variance $\sigma$ are encoded simultaneously, hence we split the network output into separate latent vectors in order to apply reparameterisation. The resulting latent representations of the images are finally subtracted using $z = z_{varying} - z_{pivot}$ to obtain the latent transformation representation. #### Decoder <center> <img src="https://hackmd.io/_uploads/Hy53dDlAJx.png" width="700"> </center> Figure 3: Visualisation of the Decoder module. It contains the Affine module, and Flow net which is a sequence of Flow modules consisting of ActNorm, invertible 1x1 conv and additive coupling layers. The ReLU non-linearities in the Additive Coupling layers are omitted for visualisation purposes.  The decoder is implemented using a normalizing flow network $f_{\theta}$, parameterized by learnable parameters $\theta$, called GLOW [[1]](##References). In essence, the latent representation $z$, modelled as a vector sampled from a multivariate Gaussian, simultaneously with the pivot image are passed through a sequence of learned transformations to obtain a transformed image. The affine layer models $\rho_{X'}(z)$ and reconstructs the latent vector $z \in \mathbb{R}^{d}$ to an image in $\mathbb{R^{C \times H \times W}}$. The output of the affine layer is $x' = x_{1} + Mz$, where M is a learnable matrix. We can think of $Mz$ as containing the necessary information needed for the normalizing flow network to apply the transformation onto the pivot image. The decoder consists of three flow nets, which contain the layers responsible for applying the invertible transformations. Each flow module contains an Activation Normalisation (ActNorm), an Invertible 1x1 Convolution and an Affine Coupling layer. In practice each net is followed by an unsqueeze layer to keep the dimensions of the latent image representation the same. The ActNorm performs an affine transformation of the activations using a scale and bias parameter per channel, which has similar stability characteristics as batch normalisation. The Invertible 1x1 Convolution applies an invertible 1 x 1 convolution over channels. The additive coupling layer is a straightforward, invertible transformation. The input $x$ is split into $x_{a}$ and $x_{b}$. A neural network $f_{\theta}$ takes $x_b$ as input. Then, $x_b$ is updated using $x_b = x_a + f_{\theta}(x_{b})$ and the output is concatenated with $x_{a}$ to obtain $x'= concat(x_{a}, x_{b})$. ### Training and Evaluation #### Dataset We are using dSprites dataset for this project [[13]](##References). DSprites is an generated dataset containing 5 latent factors, namely: - shape - scale - orientation - x-position - y-position The recombination task can be dividied into two catagories: recombination-to-element and recombination-to-range. The first setting rescricts all dimensions of the test set, here it's $shape = ellipse, pos_x \ge 0.6, pos_y \ge 0.6, 120^\circ \le orientation \le 240^\circ, scale < 0.6$. The second setting only rescricts a part of the dimensions. In contrast with the paper, we excluded $shape = ellipsis, pos_x \ge 0.5$ from training set for the to-range setting. Then the test set of to-element is a subset of that for to-range, thus we can compare the performance of different settings on the same test set.  #### Loss function  As the goal of the network is to reconstruct the transformation of an image, the model aims to minimise the binary cross entropy loss $l_{X}$ between the reconstructed image and the varying image $x_{varying}$: $$L_{recon} = l_{X}(D(E(x_{pivot}, x_{varying}),x_{pivot}), x_{varying}).$$ To enhance disentanglement, the latent output $z = E(x_{pivot}, x_{varying})$ is optimised by minimising the Kullback-Leibler divergence between the learned latent $z$ and the normal isotropic Gaussian prior $q(x)$: $$L_{KL} = D_{KL}(p(x) \lVert q(z)).$$ Technically, all the information that is required to generate a reconstruction of $x_{varying}$ is in the image $x_{varying}$. Hence, a natural tendency of the model would be to ignore $x_{pivot}$ and learn a reconstruction from only the $x_{varying}$. The latent $z$ would not actually encode a transformation, but merely a description of the image $x_{varying}$. The latent reconstruction loss puts an additional constraint on the learned transformation representation $z$: $$L_{recon\_latent}= \lVert E(x_{pivot}, D(z, x_{pivot})) - z\lVert_{1}$$ that ensures that the latent represents the group element to transform $x_{pivot}$ to $D(z, x_{pivot})$. The total loss is defined as: $$L_{MAGA}= L_{recon} + \beta_{KL}L_{KL} + \beta_{recon\_latent}L_{recon\_latent}$$ where $\beta_{KL}$ and $\beta_{recon\_latent}$ are hyperparameters that control the level of regularisation.  ## Results ### Recombination losses of reimplementation Except the BCE loss between ground truth and regenerated images under to-element and to-range settings, we also calculate that between the two settings. i.e. how the regenerated image from to-element and to-range settings different with each other. | Model | to-Element | to-Range | to-Diff | -------- | -------- | -------- | -------- | | VAE | 16.85 | 55.42 | 108.82 | | MAGANet | 18.54 | 38.62 | 68.86 | The visualization aligns with the loss that VAE performs better on to element setting, while MAGA denegrade less when trained on to range setting. <center> <img src="https://hackmd.io/_uploads/HJts28g0yg.png" width="310"> <img src="https://hackmd.io/_uploads/S1D23Ix0yl.png" width="310"> </center> Figure 4: Visualisation of recombined images from MAGA and VAE under to-element and to-range settings. ### Comparing VAE and MAGA encoding To investigate how the two models learns the representation, we visualized the encoded latent space. We used UMAP as our dimensionality reduction algorithm. Compared to t-SNE, it's much faster and scales to our dataset, and compared to PCA, it retains the semantic meaning after reduction. For all our UMAP visualisations we set the number of neighbors to 50 and number of components to 3. >You can find an interactive 3D plot at the [end of the blog post](##Conclusions). #### Using color to encode shape The clear boundaries across shapes in the MAGA visualization suggest that MAGANet is better at differentiating the shapes into the latent space than the VAE implementation. <center> <img src="https://hackmd.io/_uploads/B18PSwxCkg.png" width="800"> </center> Figure 5: Visualization of latent spaces for MAGANET(left) and VAE(right) colored by shape in varying image for MAGA and input image for VAE. #### Using color to encode position We can observe a smooth color transition when traversing the shape from one end to the other. This indicates that similar positions were encoded to adjacent positions in an ordered manner in the latent space. Although a similar ordering can be observed from the VAE implmentation, it is noisier and the transition isn't as smooth. <center> <img src="https://hackmd.io/_uploads/rkCuBNg0Jl.png" width="800"> Figure 6: Visualization of latent spaces for MAGANET(left) and VAE(right) colored by x-position in varying image for MAGA and input image for VAE. <img src="https://hackmd.io/_uploads/HkWhINx0ye.png" width="800"> </center> Figure 7: Visualization of latent spaces for MAGANET(left) and VAE(right) colored by y-position in varying image for MAGA and input image for VAE. #### Using color to encode scale MAGA exhibits a pattern with clearly defined layers for different scales. For the VAE visualisation, a similar but much weaker pattern was observed. <center> <img src="https://hackmd.io/_uploads/SJ14uVxRye.png" width="800"> </center> Figure 8: Visualization of latent spaces for MAGANET(left) and VAE(right) colored by scale in varying image for MAGA and input image for VAE. #### Using color to encode orientation An interesting result is observed when using color to encode orientation. No pattern can be observed for either the MAGA or the VAE implementation. We suspect this has to do with the fact that the shapes present in the DSprites Dataset have multiple axes of symmetry. This makes it difficult to determine the absolute orientation of these shapes. Thus leading to the apparent random organisation of the points in the plot. <center> <img src="https://hackmd.io/_uploads/H13fhEeAkg.png" width="800"> </center> Figure 9: Visualization of latent spaces for MAGANET(left) and VAE(right) colored by orientation in varying image for MAGA and input image for VAE. ### Understanding disentanglement of the latent transformation In the section above we visualised the latent space and found that the encoder maps shape information to separate parts of the latent space. This raises the question whether we can manually control properties, such as shape, of the generated image by traversing the latent space. We encode the pivot and varying images into a latent z, then fix the latent values for all but one component, which we traverse in an interval $\alpha_i \in \{\alpha_{0} - \delta, \alpha_{0} + \delta \}$, where $\alpha_{0}$ is the original value of the latent component and $\delta$ is the modification constant ($\delta = 2)$. <center> <img src="https://hackmd.io/_uploads/rkyjCdS6yl.gif" width="4200"> <img src="https://hackmd.io/_uploads/ryVoCdr6kl.gif" width="4200"> <img src="https://hackmd.io/_uploads/HJWsROHpyl.gif" width="4200"> </center> Figure 10: Visualisation of generated images. Left: pivot image. Middle: generated from modified latent representation, each column traverse one element of latent space. Right: generated from ground truth latent representation. From above visualisations we can discover some, al beit weak, disentanglement patterns. For example, traversing the latent space of the reconstructed oval (middle) we notice that traversing can result in generating a square image. Interestingly, positioning of the object does not seem to be highly influenced by a single latent component, indicating high entanglement. </figure>   ### Looking into MAGANet decoder #### The identity property A valid group action contains at least one element that can be applied without modifying the input in any way. Given that the encoders task is to encode the transformation to go from the pivot image to the varying image, it is natural to ask what happens if we set these equal. In other words, what happens if the latent variable $z$ does not contain any transformation information? We simulated the scenario where $x_{pivot} = x_{image}$ by setting $z = \vec{0}$. Interestingly, the reconstructed image is degenerate, which could be explained by the fact that the model never attempts to reconstruct its own image during training. <center class="half"> <img src="https://hackmd.io/_uploads/HJAzJVgTyl.png" width=200/> <img src="https://hackmd.io/_uploads/HyiHk4e6yl.png" width=200/> <img src="https://hackmd.io/_uploads/HJJCJVg6Je.png" width=200/> </center> Figure 11: Visualisation of identity property. Left: pivot image. Middle: varying image, same as the pivot image. Right: generated image from latent representation.  ## Conclusions The aim of this blog post was to obtain a better understanding of the original research. We implemented the proposed framework *Modelling a Group Action* for this purpose. Although our reproduction results can not directly be compared with the original results due to insufficient reproduction details, we did similarly find that MAGANet achieves greater combinatorial generalisation capability on the Recombination-2-Range evaluation datasets. Furthermore, through visualisation of the latent space we found that encoding data transformations is more separable in latent space as compared to encoding data distributions, which were generaly more noisy. However, this does not apply to the scale and orientation properties. Due to the interesting properties that Normalizing Flow networks provide it would be interesting to see whether this approach scales on larger, and importantly, more sophisticated datasets. <iframe width="100%" height="500" src="https://ignacioquintanilla99.github.io/" frameborder="1",style="background: #FFFFFF;" ></iframe> ## References [1] I. I. Vankov and J. S. Bowers, ‘Training neural networks to encode symbols enables combinatorial generalization’, Philosophical Transactions of the Royal Society B, vol. 375, no. 1791, p. 20190309, 2020.  [2] Locatello, F., Bauer, S., Lucic, M., Raetsch, G., Gelly, S., Schölkopf, B., & Bachem, O. (2019, May). Challenging common assumptions in the unsupervised learning of disentangled representations. In international conference on machine learning (pp. 4114-4124). PMLR.  [3] Kocaoglu, Murat, et al. "Causalgan: Learning causal implicit generative models with adversarial training." arXiv preprint arXiv:1709.02023 (2017).  [4]Sauer, Axel, and Andreas Geiger. "Counterfactual generative networks." arXiv preprint arXiv:2101.06046 (2021).  [5] Thiagarajan, Jayaraman, et al. "Designing counterfactual generators using deep model inversion." Advances in Neural Information Processing Systems 34 (2021): 16873-16884.  [6] Feng, Ruili, et al. "Principled knowledge extrapolation with GANs." International Conference on Machine Learning. PMLR, 2022.  [7] Hwang, G., Choi, J., Cho, H., & Kang, M. (2023, July). MAGANet: Achieving combinatorial generalization by modeling a group action. In International Conference on Machine Learning (pp. 14237-14248). PMLR. [8] Lang, S. Algebra, volume 211. Springer Science & Business  [9] Goodfellow I., Bengio Y., Courville A. (2016) Deep Learning, Chapter 19: Inference.  [10] Kingma D. P., Welling M. (2014). Auto-Encoding Variational Bayes. In 2nd International Conference on Learning Representations (ICLR).  [11] Dinh, Laurent, David Krueger, and Yoshua Bengio. "Nice: Non-linear independent components estimation." arXiv preprint arXiv:1410.8516 (2014).  [12] Kingma, Durk P., and Prafulla Dhariwal. "Glow: Generative flow with invertible 1x1 convolutions." Advances in neural information processing systems 31 (2018).  [13] L. Matthey, I. Higgins, D. Hassabis, and A. Lerchner, ‘dSprites: Disentanglement testing Sprites dataset’, 2017. [Online]. Available: https://github.com/deepmind/dsprites-dataset/. [14] Burgess, C. P., Higgins, I., Pal, A., Matthey, L., Watters, N., Desjardins, G., & Lerchner, A. (2018). Understanding disentangling in $\beta$-VAE. arXiv preprint arXiv:1804.03599.   [15] Quessard, R., Barrett, T. D., & Clements, W. R. (2020). Learning group structure and disentangled representations of dynamical environments. arXiv preprint arXiv:2002.06991.  [16] Yang, T., Ren, X., Wang, Y., Zeng, W., & Zheng, N. (2021). Towards building a group-based unsupervised representation disentanglement framework. arXiv preprint arXiv:2102.10303.  # Contributions Ignacio: Writing blogpost, VAE implementation & training Rody: Writing blogpost, repository setup, codebase setup, latent visualization, MAGANet training Zi: Writing blogpost, latent visualization, MAGANet implementation & training

Syntax	Example	Reference
# Header	Header	基本排版
- Unordered List	Unordered List
1. Ordered List	Ordered List
- [ ] Todo List	Todo List
> Blockquote	Blockquote
Bold font	Bold font
Italics font	Italics font
~~Strikethrough~~	~~Strikethrough~~
19^th^	19^th
H~2~O	H₂O
++Inserted text++	Inserted text
==Marked text==	Marked text
[link text](https:// "title")	Link
![image alt](https:// "title")	Image
`Code`	`Code`	在筆記中貼入程式碼
```javascript var i = 0; ```	`var i = 0;`
:smile:		Emoji list
{%youtube youtube_id %}	Externals
$L^aT_eX$	L^aT_eX
:::info This is a alert area. :::	This is a alert area.