FactorVAEs and Total Correlation VAEs (TCVAEs)

Beta-VAEs encourage disentanglement by increasing the weight (eta > 1) on the $KL(q_\phi(z|x) || p(z))$ term in the ELBO. However, they do so by uniformly penalizing all aspects of this KL divergence. This can sometimes lead to an undesirable trade-off, potentially over-penalizing terms that are important for learning informative representations or by excessively simplifying the posterior. FactorVAEs and Total Correlation VAEs (TCVAEs) offer more targeted approaches. These aim to directly address one of the primary statistical properties associated with disentangled representations: the independence of the latent factors. This is achieved by focusing on a quantity known as Total Correlation.

Understanding Total Correlation (TC)

At the foundation of both FactorVAE and TCVAE lies the concept of Total Correlation (TC). For a set of random variables $z = (z_1, z_2, ..., z_D)$ with a joint distribution $q(z)$ and marginal distributions $q(z_j)$, Total Correlation is defined as the Kullback-Leibler divergence between the joint distribution and the product of its marginals:

$$TC(z) = KL(q(z) || \prod_{j=1}^D q(z_j))$$

In the context of VAEs, $q(z)$ typically refers to the aggregated posterior distribution, $q(z) = \int q_\phi(z|x) p_{data}(x) dx$. This distribution represents the overall distribution of latent codes produced by the encoder when processing the entire dataset.

Essentially, TC measures the amount of statistical dependence among the variables $z_j$.

• If the latent variables $z_j$ are perfectly independent, then $q(z) = \prod_j q(z_j)$, and $TC(z) = 0$.
• If there are dependencies among the $z_j$'s, then $q(z)$ will differ from $\prod_j q(z_j)$, and $TC(z) > 0$.

The goal of disentanglement is to have each latent dimension $z_j$ correspond to a single, independent factor of variation in the data. Minimizing $TC(z)$ directly encourages these latent dimensions to become statistically independent, which is a strong proxy for disentanglement.

Total Correlation (TC) quantifies the degree of statistical dependency among the dimensions of the latent space $z$. A lower TC (right) suggests that the latent factors $z_j$ are more independent, which is a desirable property for disentangled representations. For instance, shape, color, and size might be learned as independent factors.

FactorVAE: Penalizing Total Correlation with a Discriminator

FactorVAE introduces a direct penalty for Total Correlation into the VAE objective function. The objective for FactorVAE is:

$$L_{FactorVAE} = E_{p_{data}(x), q_\phi(z|x)}[\log p_\theta(x|z)] - KL(q_\phi(z|x) || p(z)) - \lambda \cdot TC(q(z))$$

Here, $KL(q_\phi(z|x) || p(z))$ is the standard VAE KL divergence term that regularizes the per-sample posterior. The new term, $\lambda \cdot TC(q(z))$, explicitly penalizes the Total Correlation of the aggregated posterior $q(z)$, with $\lambda$ being a hyperparameter controlling the strength of this penalty.

A significant challenge is that $TC(q(z))$ is intractable to compute directly because $q(z)$ itself is intractable (it's an integral over the data distribution). FactorVAE proposes an ingenious solution: estimate $TC(q(z))$ using a discriminator network, $D(z)$.

The process involves:

1. Sampling for the Discriminator:
   • Samples from $q(z)$: These are obtained by first sampling a data point $x \sim p_{data}(x)$ (i.e., from your training batch) and then sampling $z \sim q_\phi(z|x)$.
   • Samples from $\prod_j q(z_j)$: These are generated by first sampling a batch of $z$'s from $q(z)$ as above. Then, for each dimension $j$, the values $z_j$ are permuted randomly across the samples in the batch. This breaks dependencies between dimensions while preserving the marginal distributions $q(z_j)$.
2. Training the Discriminator: The discriminator $D(z)$ is trained to distinguish between samples drawn from $q(z)$ (labeled as "real" or class 1) and samples drawn from $\prod_j q(z_j)$ (labeled as "fake" or class 0). The discriminator's objective is typically a binary cross-entropy loss: $$L_D = - E_{z \sim q(z)}[\log D(z)] - E_{z' \sim \prod_j q(z_j)}[\log(1-D(z'))]$$
3. Estimating TC: Once the discriminator is trained, the TC term can be approximated using its outputs. One common approximation is: $$TC(q(z)) \approx E_{z \sim q(z)}[\log D(z) - \log(1-D(z))]$$ This term (or a variation) is then added to the VAE's loss function (with coefficient $\lambda$) and minimized along with the other VAE components (reconstruction error and per-sample KL divergence).

By training the VAE encoder/decoder and the TC discriminator alternately, FactorVAE encourages the encoder to produce latent codes where dimensions are independent, as this makes it harder for the discriminator to distinguish $q(z)$ from $\prod_j q(z_j)$.

Total Correlation VAE ($\beta$-TCVAE): Isolating Disentanglement Factors

The $\beta$-TCVAE, proposed by Chen et al. (2018) in "Isolating Sources of Disentanglement in VAEs," takes a different route by first decomposing the average KL divergence term from the standard VAE objective, $E_{p_{data}(x)}[KL(q_\phi(z|x) || p(z))]$. Assuming a factorized prior $p(z) = \prod_j p(z_j)$ (e.g., an isotropic Gaussian $\mathcal{N}(0, I)$), this term can be broken down into three meaningful components:

1. Index-Code Mutual Information $I(x;z)$: This is $E_{p_{data}(x)}[KL(q_\phi(z|x) || q_\phi(z))]$. It measures the mutual information between the input data $x$ and the latent code $z$. A higher value means $z$ retains more information about $x$.
2. Total Correlation $TC(z)$: This is $KL(q_\phi(z) || \prod_j q_\phi(z_j))$, the same TC term discussed before, measuring dependencies in the aggregated posterior. Lower is better for disentanglement.
3. Dimension-wise KL $\sum_j KL(q_\phi(z_j) || p(z_j))$: This term encourages the marginal distribution of each latent dimension $q_\phi(z_j)$ (from the aggregated posterior) to match the corresponding marginal of the prior $p(z_j)$.

So, the average KL divergence can be written as:

$$E_{p_{data}(x)}[KL(q_\phi(z|x) || p(z))] = I(x;z) + TC(z) + \sum_j KL(q_\phi(z_j) || p(z_j))$$

A standard $\beta$-VAE penalizes all three of these terms equally with the factor $\beta$. The insight behind $\beta$-TCVAE is that for better disentanglement, we might want to specifically upweight the penalty on $TC(z)$ without necessarily increasing the penalty on $I(x;z)$ (which could hurt reconstruction) or the dimension-wise KL terms as much.

The $\beta$-TCVAE objective function modifies the ELBO to allow for different weights on these components:

$$L_{\beta-TCVAE} = E_{p_{data}(x), q_\phi(z|x)}[\log p_\theta(x|z)] - w_I \cdot I(x;z) - w_{TC} \cdot TC(z) - w_{DKL} \cdot \sum_j KL(q_\phi(z_j) || p(z_j))$$

Typically, $w_I$ (weight for mutual information) is kept at 1. The main focus is on $w_{TC}$ (often denoted as $\beta$ in this context, hence $\beta$-TCVAE), which is set greater than 1 to emphasize the minimization of Total Correlation. $w_{DKL}$ might also be adjusted.

Estimating Terms in $\beta$-TCVAE: Unlike FactorVAE, $\beta$-TCVAE usually estimates these terms, including $TC(z)$, using minibatch Monte Carlo methods without an auxiliary discriminator. For a given minibatch of data points $\{x_1, ..., x_M\}$ and their corresponding latent samples $\{z_1, ..., z_M\}$ where $z_i \sim q_\phi(z|x_i)$:

• $q_\phi(z)$ is approximated by the empirical distribution of samples $z_i$ in the minibatch.
• $q_\phi(z_j)$ (the marginals) are approximated from these samples.
• The densities like $q_\phi(z_i)$ and $q_\phi(z_{ij})$ (for the j-th component of $z_i$) within the TC formula $E_{q_\phi(z)}[\log \frac{q_\phi(z)}{\prod_j q_\phi(z_j)}]$ are estimated. For instance, $\log q_\phi(z_i)$ can be approximated using a kernel density estimator or, more commonly in practice, by leveraging the Gaussian form of $q_\phi(z|x_k)$ and averaging: $\log q_\phi(z_i) \approx \log \frac{1}{M} \sum_{k=1}^M q_\phi(z_i | x_k)$. (Note: $q_\phi(z_i|x_k)$ means evaluating the density of $z_i$ under the posterior parameterized by $x_k$). The precise estimators can be somewhat complex but are based on operations over the samples in the current minibatch.

Comparing Approaches and Practical Notes

Feature	$\beta$-VAE	FactorVAE	$\beta$-TCVAE
TC Control	Indirect, through overall KL penalty	Direct, via explicit TC term	Direct, via decomposed KL and TC term weighting
TC Estimation	Not explicitly estimated	Discriminator-based	Minibatch Monte Carlo estimation
Complexity	Simple (one hyperparameter $\beta$)	Higher (train VAE + discriminator)	Moderate (estimators can be complex)
Hyperparams	$\beta$	$\lambda$, discriminator architecture/training	$w_I, w_{TC}, w_{DKL}$
Stability	Generally stable	Can be less stable due to GAN-like training	Estimation noise can affect stability

Considerations:

• Effectiveness of TC Penalty: Both FactorVAE and $\beta$-TCVAE generally show stronger disentanglement results than $\beta$-VAE at similar levels of reconstruction quality, precisely because they target the TC more directly. $\beta$-VAE's stronger penalty on $KL(q(z|x)||p(z))$ can lead to "over-pruning" of the latent space, potentially discarding useful information for reconstruction (reducing $I(x;z)$ too much) or forcing $q(z|x)$ to be too close to $p(z)$ too aggressively.
• Estimation Challenges:
  • For FactorVAE, training the discriminator effectively is important. The quality of TC estimation depends on how well the discriminator approximates the density ratio.
  • For $\beta$-TCVAE, minibatch-based estimation of TC and other terms can be noisy, especially with small batch sizes. The accuracy of these estimates impacts the learning dynamics.
• Hyperparameter Tuning: Both FactorVAE ($\lambda$) and $\beta$-TCVAE ($w_{TC}$, etc.) introduce new hyperparameters that require careful tuning. The optimal values can be dataset-dependent.
• Computational Cost: FactorVAE adds the cost of training and running the discriminator. $\beta$-TCVAE adds computational overhead for estimating the decomposed KL terms per batch.

FactorVAEs and TCVAEs represent significant steps towards achieving more disentangled representations by moving past simple KL weighting. They provide a more grounded framework by identifying Total Correlation as a critical property to control. While they introduce their own set of complexities in terms of estimation and hyperparameter tuning, their ability to isolate and penalize sources of entanglement often leads to more interpretable and useful latent spaces, which is a primary goal in representation learning. As you implement and experiment with these models, pay close attention to the TC estimation process and the impact of the respective hyperparameters on both disentanglement metrics and reconstruction quality.