Understanding disentangling in $β$-VAE

Introduction

Representation learning lies at the core of machine learning research. From the hand-crafted feature engineering prevalent in the past to implicit representation learning of the modern deep learning approaches , it is a common theme that the performance of algorithms is critically dependent on the nature of their input representations. Despite the recent successes of the deep learning approaches , they are still far from the generality and robustness of biological intelligence . Hence, the implicit representations learnt by these approaches through supervised or reward-based signals appear to overfit to the training task and lack the properties necessary for knowledge transfer and generalisation outside of the training data distribution.

Different ways to overcome these shortcomings have been proposed in the past, such as auxiliary tasks and data augmentation . Another less explored but potentially more promising approach might be to use task-agnostic unsupervised learning to learn features that capture properties necessary for good performance on a variety of tasks . In particular, it has been argued that disentangled representations might be helpful .

A disentangled representation can be defined as one where single latent units are sensitive to changes in single generative factors, while being relatively invariant to changes in other factors . For example, a model trained on a dataset of 3D objects might learn independent latent units sensitive to single independent data generative factors, such as object identity, position, scale, lighting or colour, similar to an inverse graphics model . A disentangled representation is therefore factorised and often interpretable, whereby different independent latent units learn to encode different independent ground-truth generative factors of variation in the data.

Most initial attempts to learn disentangled representations required supervised knowledge of the data generative factors . This, however, is unrealistic in most real world scenarios. A number of purely unsupervised approaches to disentangled factor learning have been proposed , including $\beta$-VAE , the focus of this text.

$\beta$-VAE is a state of the art model for unsupervised visual disentangled representation learning. It is a modification of the Variational Autoencoder (VAE) objective, a generative approach that aims to learn the joint distribution of images $\mathbf{x}$ and their latent generative factors $\mathbf{z}$. $\beta$-VAE adds an extra hyperparameter $\beta$ to the VAE objective, which constricts the effective encoding capacity of the latent bottleneck and encourages the latent representation to be more factorised. The disentangled representations learnt by $\beta$-VAE have been shown to be important for learning a hierarchy of abstract visual concepts conducive of imagination and for improving transfer performance of reinforcement learning policies, including simulation to reality transfer in robotics . Given the promising results demonstrating the general usefulness of disentangled representations, it is desirable to get a better theoretical understanding of how $\beta$-VAE works as it may help to scale disentangled factor learning to more complex datasets. In particular, it is currently unknown what causes the factorised representations learnt by $\beta$-VAE to be axis aligned with the human intuition of the data generative factors compared to the standard VAE . Furthermore, $\beta$-VAE has other limitations, such as worse reconstruction fidelity compared to the standard VAE. This is caused by a trade-off introduced by the modified training objective that punishes reconstruction quality in order to encourage disentanglement within the latent representations. This paper attempts to shed light on the question of why $\beta$-VAE disentangles, and to use the new insights to suggest practical improvements to the $\beta$-VAE framework to overcome the reconstruction-disentanglement trade-off.

We first discuss the VAE and $\beta$-VAE frameworks in more detail, before introducing our insights into why reducing the capacity of the information bottleneck using the $\beta$ hyperparameter in the $\beta$-VAE objective might be conducive to learning disentangled representations. We then propose an extension to $\beta$-VAE motivated by these insights that involves relaxing the information bottleneck during training enabling it to achieve more robust disentangling and better reconstruction accuracy.

Variational Autoencoder (VAE)

Suppose we have a dataset $\mathbf{x}$ of samples from a distribution parametrised by ground truth generative factors $\mathbf{z}$. The variational autoencoder (VAE) aims to learn the marginal likelihood of the data in such a generative process:

where $\phi$, $\theta$ parametrise the distributions of the VAE encoder and the decoder respectively. This can be re-written as:

where D_{KL}\big{(}\;\|\;\big{)} stands for the non-negative Kullback–Leibler divergence between the true and the approximate posterior. Hence, maximising $\mathcal{L}(\theta,\phi;\mathbf{x},\mathbf{z})$ is equivalent to maximising the lower bound to the true objective in Eq. 1:

In order to make the optimisation of the objective in Eq. 3 tractable in practice, assumptions are commonly made. The prior $p(\mathbf{z})$ and posterior $q_{\phi}(\mathbf{z}|\mathbf{x})$ distributions are parametrised as Gaussians with a diagonal covariance matrix; the prior is typically set to the isotropic unit Gaussian $\mathcal{N}(0,1)$. Parametrising the distributions in this way allows for use of the “reparametrisation trick” to estimate gradients of the lower bound with respect to the parameters $\phi$, where each random variable $z_{i}\sim q_{\phi}(z_{i}|\mathbf{x})=\mathcal{N}(\mu_{i},\sigma_{i})$ is parametrised as a differentiable transformation of a noise variable $\epsilon\sim\mathcal{N}(0,1)$:

β𝛽\beta-VAE

$\beta$-VAE is a modification of the variational autoencoder (VAE) framework that introduces an adjustable hyperparameter $\beta$ to the original VAE objective:

Well chosen values of $\beta$ (usually $\beta>1$) result in more disentangled latent representations $\mathbf{z}$. When $\beta=1$, the $\beta$-VAE becomes equivalent to the original VAE framework. It was suggested that the stronger pressure for the posterior $q_{\phi}(\mathbf{z}|\mathbf{x})$ to match the factorised unit Gaussian prior $p(\mathbf{z})$ introduced by the $\beta$-VAE objective puts extra constraints on the implicit capacity of the latent bottleneck $\mathbf{z}$ and extra pressures for it to be factorised while still being sufficient to reconstruct the data $\mathbf{x}$ . Higher values of $\beta$ necessary to encourage disentangling often lead to a trade-off between the fidelity of $\beta$-VAE reconstructions and the disentangled nature of its latent code $\mathbf{z}$ (see Fig. 6 in ). This due to the loss of information as it passes through the restricted capacity latent bottleneck $\mathbf{z}$.

Understanding disentangling in β𝛽\beta-VAE

The $\beta$-VAE objective is closely related to the information bottleneck principle :

where $I(\cdot;\cdot)$ stands for mutual information and $\beta$ is a Lagrange multiplier. The information bottleneck describes a constrained optimisation objective where the goal is to maximise the mutual information between the latent bottleneck $Z$ and the task $Y$ while discarding all the irrelevant information about $Y$ that might be present in the input $X$. In the information bottleneck literature, $Y$ would typically stand for a classification task, however the formulation can be related to the auto-encoding objective too .

2 β𝛽\beta-VAE through the information bottleneck perspective

In $\beta$-VAE, the posterior $q(\mathbf{z}|\mathbf{x})$ is encouraged to match the unit Gaussian prior $p(z_{i})=\mathcal{N}(0,1)$. Since the posterior and the prior are factorised (i.e. have diagonal covariance matrix) and posterior samples are obtained using the reparametrization (Eq. 4) of adding scaled independent Gaussian noise $\sigma_{i}\epsilon_{i}$ to a deterministic encoder mean $\mu_{i}$ for each latent unit $z_{i}$, we can take an information theoretic perspective and think of $q(\mathbf{z}|\mathbf{x})$ as a set of independent additive white Gaussian noise channels $z_{i}$, each noisily transmitting information about the data inputs $x_{n}$. In this perspective, the KL divergence term D_{KL}\big{(}q_{\phi}(\mathbf{z}|\mathbf{x})\;\|\;p(\mathbf{z})\big{)} of the $\beta$-VAE objective (see Eq. 5) can be seen as an upper bound on the amount of information that can be transmitted through the latent channels per data sample (since it is taken in expectation across the data). The KL divergence is zero when $q(z_{i}|\mathbf{x})=p(\mathbf{z})$, i.e $\mu_{i}$ is always zero, and $\sigma_{i}$ always 1, meaning the latent channels $z_{i}$ have zero capacity. The capacity of the latent channels can only be increased by dispersing the posterior means across the data points, or decreasing the posterior variances, which both increase the KL divergence term.

3 Comparing disentangling in β𝛽\beta-VAE and VAE

A representation learned under a weak bottleneck pressure (as in a standard VAE) can exhibit this locality property in an incomplete, fragmented way. To illustrate this, we trained a standard VAE (i.e. with $\beta=1$) and a $\beta$-VAE on a simple dataset with two generative factors of variation: the $x$ and $y$ position of a Gaussian blob (Fig. 2). The standard VAE learns to represent these two factors across four latent dimensions, whereas $\beta$-VAE represents them in two. We examine the nature of the learnt latent space by plotting its traversals in Fig. 2, whereby we first infer the posterior $q(\mathbf{z}|\mathbf{x})$, before plotting the reconstructions resulting from modifying the value of each latent unit $z_{i}$ one at a time in the $$range while keeping all the other latents fixed to their inferred values. We can see that the$\beta$-VAE represention exhibits the locality property described in Sec. 4.2 since small steps in each of the two learnt directions in the latent space result in small changes in the reconstructions. The VAE represention, however, exhibits fragmentation in this locality property. Across much of the latent space, small traversals produce reconstructions with small, consistent offsets in the position of the sprite, similar to$\beta$-VAE. However, there are noticeable representational discontinuities, at which small latent perturbations produce reconstructions with large or inconsistent position offsets. Reconstructions near these boundaries are often of poor quality or have artefacts such as two sprites in the scene.

We have seen how a strong pressure for overlapping posteriors encourages $\beta$-VAE to find a representation space preserving as much as possible the locality of points on the data manifold. However, why would it find representational axes that are aligned with the generative factors of variation in the data? Our key hypothesis is that $\beta$-VAE finds latent components which make different contributions to the log-likelihood term of the cost function (Eq. 5). These latent components tend to correspond to features in the data that are intuitively qualitatively different, and therefore may align with the generative factors in the data.

At this point we can ask what pressures could encourage this new factor of variation to be encoded into a distinct latent dimension. We hypothesise that two properties of $\beta$-VAE encourage this. Firstly, embedding this new axis of variation of the data into a distinct latent dimension is a natural way to satisfy the data locality pressure described in Sec. 4.2. A smooth representation of the new factor will allow an optimal packing of the posteriors in the new latent dimension, without affecting the other latent dimensions. We note that this pressure alone would not discourage the representational axes from rotating relative to the factors. However, given the differing contributions each factor makes to the reconstruction log-likelihood, the model will try to allocate appropriately differing average capacities to the encoding axes of each factor (e.g. by optimising the posterior variances). But, the diagonal covariance of the posterior distribution restricts the model to doing this in different latent dimensions, giving us the second pressure, encouraging the latent dimensions to align with the factors.

We tested these intuitions by training a simplified model to generate dSprites conditioned on the ground-truth factors, $\mathbf{f}$, with a controllable information bottleneck (Fig. 3). In particular, we wanted to evaluate how much information the model would choose to retain about each factor in order to best reconstruct the corresponding images given a total capacity constraint. In this model, the factors are each independently scaled by a learnable parameter, and are subject to independently scaled additive noise (also learned), similar to the reparameterised latent distribution in $\beta$-VAE. This enables us to form a KL divergence of this factor distribution with a unit Gaussian prior. We trained the model to reconstruct the images with samples from the factor distribution, but with a range of different target encoding capacities by pressuring the KL divergence to be at a controllable value, $C$. The training objective combined maximising the log likelihood and minimising the absolute deviation from $C$ (with a hyperparameter $\gamma$ controlling how heavily to penalise the deviation, see Sec. A.2):

In practice, a single model was trained across of range of $C$’s by linearly increasing it from a low value (0.5 nats) to a high value (25.0 nats) over the course of training (see top left panel in Fig. 3). Consistent with the intuition outlined above, at very low capacities ($C<5$ nats), the KLs for all the factors except the X and Y position factors are zero, with $C$ always shared equally among X and Y. As expected, the model reconstructions in this range are blurry, only capturing the position of the original input shapes (see the bottom row of the lower panel in Fig. 3). However, as $C$ is increased, the KLs of other factors start to increase from zero, at distinct points for each factor. For example, starting around $C=6$ nats, the KL for the scale factor begins to climb from zero, and the model reconstructions become scaled (see 7.3 nats row in lower panel of Fig. 3). This pattern continues until all factors have a non-zero KL and eventually the reconstructions begin to look almost identical to the samples.

Improving disentangling in β𝛽\beta-VAE with controlled capacity increase

The intuitive picture we have developed of gradually adding more latent encoding capacity, enabling progressively more factors of variation to be represented whilst retaining disentangling in previously learned factors, motivated us to extend $\beta$-VAE with this algorithmic principle. We applied the capacity control objective from the ground-truth generator in the previous section (Eq. 7) to $\beta$-VAE, allowing control of the encoding capacity (again, via a target KL, $C$) of the VAE’s latent bottleneck, to obtain the modified training objective:

Similar to the generator model, $C$ is gradually increased from zero to a value large enough to produce good quality reconstructions (see Sec. A.2 for more details).

Results from training with controlled capacity increase on coloured dSprites can be seen in Figure 4(a), which demonstrate very robust disentangling of all the factors of variation in the dataset and high quality reconstructions.

Single traversals of each latent dimension show changes in the output samples isolated to single data generative factors (second row onwards, with the latent dimension traversed ordered by their average KL divergence with the prior, high KL to low). For example, we can see that traversal of the latent with the largest KL produces smooth changes in the Y position of the reconstructed shape without changes in other factors. The picture is similar with traversals of the subsequent latents, with changes isolated to X position, scale, shape, rotation, then a set of three colour axes (the last two latent dimensions have an effectively zero KL, and produce no effect on the outputs).

Furthermore, the quality of the traversal images are high, and by eye, the model reconstructions (second row) are quite difficult to distinguish from the corresponding data samples used to generate them (top row). This contrasts with the results previously obtained with the fixed $\beta$-modulated KL objective in .

We also trained the same model on the 3D Chairs dataset , with latent traversals shown in Figure 4(b). We can see that reconstructions are of high quality, and traversals of the latent dimensions produce smooth changes in the output samples, with reasonable looking chairs in all cases. With this richer dataset it is unclear exactly what the disentangled axes should correspond to, however, each traversal appears to generate changes isolated in one or few qualitative features that we might identify intuitively, such as viewing angle, size, and chair leg and back styles.

Conclusion

We have developed new insights into why $\beta$-VAE learns an axis-aligned disentangled representation of the generative factors of visual data compared to the standard VAE objective. In particular, we identified pressures which encourage $\beta$-VAE to find a set of representational axes which best preserve the locality of the data points, and which are aligned with factors of variation that make distinct contributions to improving the data log likelihood. We have demonstrated that these insight produce an actionable modification to the $\beta$-VAE training regime. We proposed controlling the increase of the encoding capacity of the latent posterior during training, by allowing the average KL divergence with the prior to gradually increase from zero, rather than the fixed $\beta$-weighted KL term in the original $\beta$-VAE objective. We show that this promotes robust learning of disentangled representation combined with better reconstruction fidelity, compared to the results achieved in the original formulation of .

References

Appendix A Supplementary Materials

The neural network models used for experiments in this paper all utilised the same basic architecture. The encoder for the VAEs consisted of 4 convolutional layers, each with 32 channels, 4x4 kernels, and a stride of 2. This was followed by 2 fully connected layers, each of 256 units. The latent distribution consisted of one fully connected layer of 20 units parametrising the mean and log standard deviation of 10 Gaussian random variables (or 32 for the CelebA experiment). The decoder architecture was simply the transpose of the encoder, but with the output parametrising Bernoulli distributions over the pixels. ReLU activations were used throughout. The optimiser used was Adam with a learning rate of 5e-4.

A.2 Training Details

$\gamma$ used was 1000, which was chosen to be large enough to ensure the actual KL was always close to the target KL, $C$. For dSprites, $C$ was linearly increased from 0 to 25 nats over the course of 100,000 training iterations, for CelebA it was increased to 50 nats.