Metrics for Quantifying Disentanglement

Disentangled representations aim to map distinct generative factors in data to distinct dimensions in the latent space. Quantifying how well a model achieves this objective is a complex task. Observing generated samples or latent traversals offers qualitative observations, but precise comparison and optimization necessitate objective, quantitative measures. Several widely adopted metrics for evaluating disentanglement are presented, focusing on scenarios where ground-truth factors of variation are known.

"It's important to preface this discussion by acknowledging a common challenge: most disentanglement metrics rely on access to datasets with annotated, ground-truth factors of variation (e.g., in an image dataset, these could be object shape, color, position, scale). Datasets like dSprites, Shapes3D, and MPI3D were specifically created for this purpose. In many applications, such ground-truth factors are unavailable, making the direct application of these metrics difficult. Nevertheless, they serve as invaluable tools for research and for understanding the properties of models that encourage disentanglement."

The diagram below illustrates the ideal scenario these metrics aim to quantify.

Idealized disentangled representation. Each ground-truth factor ($g_i$) strongly influences a specific latent dimension ($z_i$), while interactions with other latent dimensions are weak.

Let's examine some of the prominent metrics used to assess this ideal.

Mutual Information Gap (MIG)

The Mutual Information Gap (MIG) score (Chen et al., 2018) aims to quantify the degree to which individual latent dimensions $z_j$ capture individual ground-truth factors $g_k$. The core idea is that for a well-disentangled representation, each factor $g_k$ should have high mutual information with one specific latent dimension and low mutual information with all others.

Calculation:

1. Prerequisite: A dataset with $N$ samples, $D$ latent dimensions $z = (z_1, ..., z_D)$, and $K$ known, discrete ground-truth factors $g = (g_1, ..., g_K)$.
2. Estimate Mutual Information: For each pair of latent dimension $z_j$ and ground-truth factor $g_k$, estimate the mutual information $I(z_j; g_k)$. If $z_j$ is continuous, it's typically discretized by binning its empirical distribution. The mutual information is then computed using the joint and marginal probability distributions over the (now discrete) $z_j$ and $g_k$. $$I(z_j; g_k) = \sum_{val(z_j)} \sum_{val(g_k)} p(z_j, g_k) \log \frac{p(z_j, g_k)}{p(z_j)p(g_k)}$$
3. Identify Top Latents per Factor: For each ground-truth factor $g_k$:
   • Find the latent dimension $z_{j_1}$ that has the highest mutual information with $g_k$: $j_1 = \text{argmax}_j I(z_j; g_k)$.
   • Find the latent dimension $z_{j_2}$ that has the second highest mutual information with $g_k$ among the remaining latents ($j_2 \neq j_1$).
4. Calculate Normalized Gap: The gap for factor $g_k$ is the difference in mutual information, normalized by the entropy of the factor $H(g_k)$: $$\text{Gap}_k = \frac{I(z_{j_1}; g_k) - I(z_{j_2}; g_k)}{H(g_k)}$$ Normalization by $H(g_k)$ accounts for factors that are inherently easier or harder to predict.
5. Average Gaps: The MIG score is the average of these normalized gaps over all $K$ ground-truth factors: $$\text{MIG} = \frac{1}{K} \sum_{k=1}^K \text{Gap}_k$$

Interpretation:

• A higher MIG score is better, indicating greater disentanglement. A score closer to 1 suggests that for each factor, one latent dimension is significantly more informative than any other.
• A score near 0 implies that factors are either not captured well, or the information is spread across multiple latent dimensions (entangled).

Strengths:

• Intuitive: Directly measures the "informativeness gap" which aligns well with the idea of one-to-one mapping.
• Normalization: Dividing by $H(g_k)$ provides a scale-invariant measure.

Weaknesses:

• Requires ground-truth factors.
• Estimation of mutual information, especially with discretization of continuous latents, can be sensitive to the number of bins and amount of data.
• Only considers the top two latents for each factor, potentially missing more complex entanglement patterns.

Separated Attribute Predictability (SAP)

The Separated Attribute Predictability (SAP) score (Kumar et al., 2017) takes a slightly different approach. Instead of directly measuring mutual information, it assesses how well each ground-truth factor can be predicted from individual latent dimensions using a simple classifier. The "separation" comes from the difference in predictability offered by the best versus the second-best latent dimension for a given factor.

Calculation:

1. Prerequisite: Similar to MIG, a dataset with latent codes $z$ and corresponding ground-truth factors $g$.
2. Train Predictors: For every pair of a latent dimension $z_j$ and a ground-truth factor $g_k$:
   • Train a straightforward classifier (e.g., logistic regression, linear Support Vector Machine) to predict the factor $g_k$ using only the values of the single latent dimension $z_j$.
   • Evaluate the classifier's performance, typically using classification accuracy. This results in a score matrix $S$, where $S_{jk}$ is the accuracy of predicting $g_k$ from $z_j$.
3. Identify Top Predicting Latents per Factor: For each ground-truth factor $g_k$:
   • Find the latent dimension $z_{j_1}$ that yields the highest prediction score for $g_k$: $S_{j_1k} = \max_j S_{jk}$.
   • Find the latent dimension $z_{j_2}$ that yields the second highest prediction score for $g_k$ from the remaining latents.
4. Calculate Score Difference: The SAP contribution for factor $g_k$ is the difference between these top two scores: $$\text{SAP}_k = S_{j_1k} - S_{j_2k}$$
5. Average Differences: The overall SAP score is the average of these differences across all $K$ ground-truth factors: $$\text{SAP} = \frac{1}{K} \sum_{k=1}^K \text{SAP}_k$$

Interpretation:

• A higher SAP score indicates better disentanglement. It suggests that for each factor, one latent dimension is substantially more predictive than any other single dimension.
• A low SAP score implies that either factors are not well predicted by individual latents, or multiple latents are similarly (and perhaps weakly) predictive.

Strengths:

• Uses a predictive framework, which can be more stable than MI estimation in some cases.
• The concept of "predictability difference" clearly targets the separation aspect.

Weaknesses:

• Requires ground-truth factors.
• The choice of classifier and its hyperparameters can influence the results. The metric's outcome might reflect the chosen classifier's ability as much as the representation's quality.
• It measures predictability rather than direct information content, which are related but not identical.

Disentanglement, Completeness, and Informativeness (DCI)

The DCI framework (Eastwood & Williams, 2018) offers a more multifaceted evaluation by proposing three distinct scores: Disentanglement, Completeness, and Informativeness. These metrics leverage the feature importances derived from training a predictor (often a Lasso regressor or a Random Forest) to predict each ground-truth factor using the entire set of latent dimensions.

1. Disentanglement Score:

   • Intuition: Measures whether each latent dimension $z_j$ focuses on encoding a small number of ground-truth factors. Ideally, a single $z_j$ should be important for only one $g_k$.
   • Calculation (Simplified):
     1. For each factor $g_k$, train a predictor (e.g., Lasso) using all latent dimensions $z=(z_1, ..., z_D)$ and extract feature importances $R_{jk}$ (importance of $z_j$ for predicting $g_k$).
     2. For each latent $z_j$, consider its importance vector across all factors: $(R_{j1}, R_{j2}, ..., R_{jK})$. Normalize this vector so its elements sum to 1, forming a probability distribution.
     3. The disentanglement score for $z_j$ is $1 - H(P_j)$, where $H(P_j)$ is the entropy of this normalized importance distribution. If $z_j$ is important for only one factor, the distribution is peaked, entropy is low, and the score is high.
     4. The overall Disentanglement score is the average of these scores over all "active" latents (those that have non-negligible importance for at least one factor), possibly weighted by how predictive each latent is overall.
   • Interpretation: A high DCI-Disentanglement score (closer to 1) indicates that individual latent dimensions tend to specialize in representing single factors of variation.

2. Completeness Score:

   • Intuition: Measures whether each ground-truth factor $g_k$ is primarily captured by a small number of latent dimensions. Ideally, a single $g_k$ should be predictable from only one $z_j$.
   • Calculation (Simplified):
     1. Using the same importance matrix $R_{jk}$ from above.
     2. For each factor $g_k$, consider its importance vector across all latents: $(R_{1k}, R_{2k}, ..., R_{Dk})$. Normalize this vector to sum to 1.
     3. The completeness score for $g_k$ is $1 - H(Q_k)$, where $H(Q_k)$ is the entropy of this normalized importance distribution. If $g_k$ is mainly explained by one latent, entropy is low, and score is high.
     4. The overall Completeness score is the average of these scores over all factors.
   • Interpretation: A high DCI-Completeness score (closer to 1) suggests that each ground-truth factor is well-represented by a small subset (ideally one) of latent dimensions, rather than its information being diffusely spread.

3. Informativeness Score:

   • Intuition: Measures how well the learned latent representation $z$ as a whole can be used to predict the ground-truth factors $g_k$. This is a basic check of whether the representation has captured useful information about the factors at all.
   • Calculation: Typically, this is the average prediction accuracy (e.g., $R^2$ for regression, classification accuracy for classification) when predicting each $g_k$ from the full latent code $z$.
   • Interpretation: A high Informativeness score means the latents are useful for downstream prediction tasks related to the factors. Low informativeness suggests the VAE failed to learn a meaningful representation of these factors.

Strengths of DCI:

• Provides a more comprehensive view by separating concerns of disentanglement, completeness, and overall utility.
• Using feature importances from capable predictors like Random Forests or Lasso can capture non-linear relationships and provide estimates.

Weaknesses of DCI:

• Requires ground-truth factors.
• The choice of predictor model (e.g., Lasso, Random Forest) and its hyperparameters can affect the importance scores and thus the final DCI metrics.
• Can be more computationally intensive and complex to implement than MIG or SAP.

General Remarks on Disentanglement Metrics

While MIG, SAP, and DCI are common, other metrics exist, each with its own details. Most share the characteristic of relying on ground-truth factors. It's also worth noting that:

• No Single Perfect Metric: Disentanglement is a multifaceted property, and no single metric perfectly captures all its aspects. Different metrics might rank models differently.
• Sensitivity to Hyperparameters: Many metrics involve choices (number of bins for MI, type of classifier for SAP/DCI) that can influence scores. Consistent evaluation protocols are important.
• Correlation with Qualitative Assessment: Ideally, quantitative metrics should correlate with human judgment of disentanglement based on latent traversals. However, this is not always straightforward.
• Focus on Axis-Alignment: Most of these metrics implicitly favor axis-aligned disentanglement, where each factor aligns with a single coordinate axis in the latent space. More general notions of disentanglement (e.g., factors encoded in linear subspaces) are harder to quantify with these standard tools.

Understanding these metrics is important for anyone working on disentangled representation learning. They provide the tools to move past qualitative hunches and to rigorously assess progress in developing VAEs and other models that learn more interpretable and structured latent spaces. As you proceed to the hands-on practical later in this chapter, you will have the opportunity to implement some of these metrics to evaluate the VAEs you train.