I am investigating methods of learning disentangled representations of data with deep neural networks and their applications in the medical domain. Deep representation learning is an ever growing field, but as the computational capabilities of modern neural networks continue to grow, the problem of their interpretability remains open. This is especially crucial in medical tasks, where working with purely "black-box" models might be simply not acceptable. By learning disentangled representations of data we aim for greater interpretability, transferability and even the ability to answer counterfactual questions.

In my recent work I investigated the relations between dataset properties and disentanglement, uncovering the role of structures of variance caused by the underlying factors. 

Problem Setting and Definition

Representation learning algorithms aim to learn transformations mapping complex, usually high-dimensional data (e.g., images) into more compact (lower dimensional) feature representations. These should ideally describe more abstract characteristics of the data (e.g., objects in an image). This can be ilustrated with the following diagram:

where \(x_1 ... x_l \) denote the dimensions of the observed data and \(r_1 ... r_m \) dimensions of the learned representation. Note that usually \(m << k \).

In the setting of learning disentangled representations we assume that the observed data were caused by a set of so called ground-truth factors of variation, via a (possibly unknown) generative process. Typically we further assume that these factors are mutually independent, and come from a given prior distribution (e.g., normal):

A learned representation is then assumed to be disentangled if each of these factors would only influence distinct sets of the learned features. More formally, we want the following to hold:

\( \forall i : | \{ j \in \{1 \ldots k\}: \frac{\partial r_i}{\partial z_j} \neq 0 \} | \leq 1 \),

where \( \frac{\partial r_i}{\partial z_j} \) denotes the derivative of the i-th learned feature wrt. the j-th ground-truth factor of variation. Reversely, a representation is fully entangled if a change in just one factor \( z_j \) would induce a change in each of the learned features \( r_i(x) \). Ideally, one would want the learned representation to match the ground-truth factors (up to a permutation of the dimensions).

Disentaglement is believed to be a crucial property when learning representations, allowing for better interpretability and transferability across tasks (Bengio et al. [1]).

The most popular approach to learning disentangled representations is to employ models from the Variational Autoencoders (VAEs) framework (Kingma, Welling [2]), which combine Variational Inference with Autoencoder models. A VAE is a stochastic generative model, consisting of two (usually symmetric) networks - the encoder and the decoder. They are trained to jointly model the distribution of the observed data \( p(x) \), with a generative model, along with a conditional distribution of latent variables \( p(z|x) \), with a predictive model. The latter is defined by the encoder network and serves as the aforementioned latent represenation (encodings) of the data.

In standard autoencoder models the encoder yields a point estimate of the latent representation \( z_i = E(x_i) \). In a VAE, on the other hand, we obtain parameters of a distribution over possible encodings \( \mu_i,\,\sigma^{2}_i = E(x_i) \), from which we then sample: \( z_i \sim \mathcal{N}(\mu_i,\,\sigma^{2}_i ) \). 

A key aspect of this framework is the regularization imposed on these latent encodings, forcing each of them to be close to a prior distribution (usually a standard normal distribution), measured with Kullback-Leibler divergence. The tradeoff between regularization and reconstruction introduces an additional informational bottleneck, which is believed to force the model to ignore less important factors of variation in the data.

Most approaches to obtaining disentanglement can be divided in two groups - ones that modify the regularization term (e.g., by increasing its corresponding weight) or by enforcing statistical independence of dimensions of the latent space representations, effectively imposing a factorising prior on the aggregated posterior of the encodings. 

However, a large-scale study (Locatello et al. [3]) showed that all these methods do not yield consistent performance (in terms of disentangling), with high variance of achieved scores and low correlation between the selected method and level of disentanglement obtained. Also, and perhaps more importantly, they introduced the so-called Impossibility Result. It states that in a purely unsupervised setting it is impossible to identify disentangled representations, even when succesfully imposing the factorising prior.

Below we can see a visual investigation of representations of two trained models - baseline (5 columns on the left) and the proposed Ada-GVAE model (5 columns on the right). Odd rows show images generated by interpolating across a single latent dimension while even rows show difference images with respect to the mean (middle) image. The Ada-GVAE (right) model was able to better disentangle intensity (top rows) from features related to aging or neurodegeneration, such as changes in ventricle or fissure sizes (bottom rows). In the baseline model (left) these features are more entangled with each other - both ventricle size and intensity are changing in each row.

As the variances differ more (left plot), the PCA solution is relatively well-aligned with the ground-truth - however as the variances become closer (right plot), the PCA solution can is more likely to become misaligned. We observed a similar effect with models trained on both toy and benchmark datasets - they seem to disentangle better when the underlying ground-truth factors of variation induce different amounts of variance in the data.

We also investigate the effect of global consistency of the principal directions of variance, i.e., whether a ground-truth factor of variance induces the same amount of variance in the data globally on the data manifold. An example of globally inconsistent data can be seen below - changing the color of the object influences a different number of pixels depending on the object's size (position on the size-corresponding axis in the latent space):

Analysing results on toy and benchmark datasets, we observe that VAEs disentangle the better on datasets with globally consistent structures of variance. Additionally, we note that such variance inconsistencies lead not only to entanglement (incorrect alignment of axes), but also to inconsistency of encodings (factors being encoded using different dimensions of the representation, depending on the position on the data manifold).

Sparsity in machine learning models can be referring to either a sparse output of the model (e.g., a feature vector with many zero entries) or to sparsity of the solution (many parameter weights equal or close to zero). 

In this project I investigate the latter - that is sparsity of neural networks - and its relation to disentanglement. Standard neural network models have dense channel-to-channel connectivity. However, it is not necessary for a feature (represented by a hidden unit) to be constructed using all available features from the subsequent layer - instead, it should suffice for it to be a combination of only a subset of them, without a loss in performance. This has been verified empirically by methods such as the Lottery Ticket Hypothesis, showing that in many cases a model with 90% less parameters can still achieve similar performance (Frankie, Carbin [4]). 

A motivation behind a possible relation between model sparsity and disentanglement can be estabilished by revision of the proof of the aforementioned impossibility result. Even if a model is powerful enough to recover all underlying factors of variation of the data and encode them in the latent space, it suffices to rotate this space to obtain an entangled representation. Because the employed prior distributions are rotationally invariant it is then impossible to identify a correct model. I argue that this entanglement can happen already in the hidden layers of a model. This is due to the dense structure of modern neural networks, since there is no incentive for the model to use the connections sparsely, especially given a random initialization of weights. Explicitly enforcing sparsity might thus alleviate this potential issue, by favoring models which "rotate the features less". Preference for sparsity can also be found in the literature regarding causality. By the Independent Causal Mechanisms Principle, the casual mechanisms (which can be modeled as hidden layers of a deep network) “should usually not affect all factors simultaneously“, that is, they should propagate down the network in a sparse and local manner (Schölkopf [5]).

A naïve approach to obtaining sparsity would be to simply employ the \( L_1 \) weight penalty. While this is plausible for 2-dimensional weight matrices, used by fully-connected layers, it does not have to have the desired effect in case of convolutional layers, whose weight tensors are at least 3-dimensional. Instead, we are interested in limiting the number of channel-to-channel connections - that is, zeroing out all weights in a convolutional filter that are connected to a particular channel of a previous layer. We implement this by introducing an additional masking tensor to each layer. Before computing a layer's output, its weight tensor is first multiplied with this mask tensor. Thus each entry in the mask effectively serves as a gate, controlling flow of information between a pair of channels from subsequent layers. This is ilustrated in the figure below:

More formally, for a convolution layer with a corresponding weight tensor \( K \in \mathbb{R}^{m \times n \times w \times h} \) we are interested in zeroing out all entries of slices \( K_{i, j, :, :} \) of shape \( w \times h \), for certain pairs of indices \( i, j \) indexing channels in the layer's input and output respectively. To achieve this, we extend the layers that we want to impose sparsity on with an additional masking tensor \( M \in \mathbb{R}^{m \times n} \). During the forward pass a masked weight tensor \( K' = K \odot M \) is used to compute the output, where \( \odot \) denotes the Hadamard product. This method also applies in the exact same way to fully-connected layers, with the only difference being that the original weight tensor is 2-dimensional.

Preliminary Results

The plot below shows results obtained on the dSprites and Cars3D datases (rows), for 3 disentanglement metrics - Mutual Information Gap (MIG), DCI Disentanglement and Modularity (columns):

For each metric/dataset combination we plot the mean results over a range of 6 regularization strengths, along with the baseline model (dashed line). For certain regularization strengths there is an improvement over the baseline. However, extremely high values seem to hinder performance for the MIG and DCI metrics, while Modularity seems to exhibit an opposite trend. Overall, there seems to be no "sweet spot" for the hyperparameter setting that would transfer well across datasets and metrics.
 

Co-tutored the Master Thesis of Dominika Matus: "Recognition of Epilepsy Seizures in EEG Data Using Deep Learning Methods"