Introduction

Variational autoencoders (VAEs) have recently been used for unsupervised disentanglement learning of complex density distributions. Numerous variants exist to encourage disentanglement in latent space while improving reconstruction. However, none have simultaneously managed the trade-off between attaining extremely low reconstruction error and a high disentanglement score.

We present a generalized framework to handle this challenge under constrained optimization and demonstrate that it outperforms state-of-the-art existing models as regards disentanglement while balancing reconstruction. We introduce three controllable Lagrangian hyperparameters to control reconstruction loss, KL divergence loss and correlation measure. We prove that maximizing information in the reconstruction network is equivalent to information maximization during amortized inference under reasonable assumptions and constraint relaxation.

This blog post focuses on the mathematical derivations of the GCVAE framework, particularly the detailed proofs from the appendix that establish the theoretical foundation of this approach.

Background: Variational Autoencoders

VAEs are a class of generative model proposed to model complex distributions existing in images, natural language and functional data. We formally define the VAE model by observing a d-dimensional input space {x_i}_i=1^N ∈ X consisting of N-independently and identically distributed (i.i.d) samples; k-dimensional latent space {z_i}_i=1^N ∈ Z (where k ≪ d) over which a generative model is defined.

We assume an empirical prior distribution p_θ(z) ∼ N(0, I) to infer an approximate posterior distribution q_φ(z|x) ∼ N(z|μ_φ(x), σ_φ²(x)I), with mean μ_φ(x) and variance σ_φ²(x)I used for reparameterization sampling of the latent space z. We model the data using conditional distribution p_θ(x|z) ∼ N(x|μ_θ(x), σ_θ²(x)I).

The objective function L(θ, φ) to be maximized is given by the Evidence Lower Bound (ELBO):

L(θ, φ) = E_{p_D}E_{z∼q_φ(z|x)}[ln p_θ(x|z)] − D_KL(q_φ(z|x)||p_θ(z))

where φ and θ represent the parameters of the neural network encoder and decoder respectively. The first term is a reconstruction error, and the second term is the KL-divergence between the approximate posterior and the prior.

The GCVAE Framework

We propose a generalized framework for variational inference modeling, taking into account the idea of information maximization in the reconstruction network. We prioritize the disentanglement of the latent space and balancing the trade-off between disentanglement metric and reconstruction loss by maximizing the mutual information between the reconstructed data x′ and latent space z.

Let I_p(x′, z) be the mutual information between the reconstructed data x′ and z under joint distribution p_θ(x′|z)p_θ(z), where p_θ(z) ∼ N(0, I). I_q(x, z) is the mutual information between x and z under joint distribution q_φ(z|x)p_D(x).

The constraint optimization formulation can be written as:

max_φ,θ I_p(x′, z)

s.t D_KL(q_φ(z|x)||p_θ(z)) ≤ ξ₁

s.t D_KL(q_φ(z)||p_θ(z)) ≤ ξ₂

ξ₁, ξ₂ ≥ 0

This implies that accurately reconstructing the original distribution p_D(x) requires us to maximize the mutual information I_p(x′, z) in the reconstructed space while reducing information loss I_q(x, z) during inference.

Figure 1: GCVAE framework. α_t, β_t and γ_t respectively provide automatic balancing of the log-likelihood and KL divergences for optimal reconstruction and disentanglement. The feed-ins A_t, B_t and C_t are expectations of variational loss.

Mathematical Derivations

The following sections present detailed mathematical derivations from the appendix of the paper, establishing the theoretical foundation of GCVAE.

Appendix A.1: Proof that Maximizing Negative KL-Divergence is Equivalent to Maximizing ELBO

We prove that maximizing the negative KL-divergence is equivalent to maximizing the ELBO.

Starting from the objective function:

L(θ, φ) = −E_{p_D}[D_KL(q_φ(z|x)||p_θ(z|x))]

Expanding the KL-divergence term:

= −E_{p_D}E_{z∼q_φ(z|x)}[ln (q_φ(z|x) / p_θ(z|x))]

Using Bayes' theorem, p_θ(z|x) = p_θ(x|z)p_θ(z) / p_θ(x), we can rewrite this as:

= −E_{p_D}E_{z∼q_φ(z|x)}[ln (q_φ(z|x) / p_θ(z) · 1 / p_θ(x|z) · p_θ(x))]

Rearranging terms, we obtain:

ln p_θ(x) ≥ E_{p_D}[ln p_θ(x|z) − D_KL(q_φ(z|x)||p_θ(z))]

This establishes that maximizing the negative KL-divergence is equivalent to maximizing the ELBO, which provides a lower bound on the log-likelihood.

Appendix A.2: Minimizing Mutual Information I_q(x, z)

We demonstrate that minimizing the mutual information between x and latent z equivalently maximizes components of the ELBO and consequently, I_p(x′, z) subject to inference constraints.

The mutual information for the inference network is expressed as:

I_q(x, z) = ∫_x∫_z q_φ(x, z) ln (q_φ(x, z) / q_φ(z)q_φ(x)) dx dz

= D_KL(q_φ(x, z)p_θ(z)||q_φ(z)q_φ(x)p_θ(z))

Expanding this expression step by step:

I_q(x, z) = E_{z∼q_φ(z|x)}[ln (q_φ(z|x) / q_φ(z) · p_θ(x, z) / p_θ(x, z))]

Using the relationship q_φ(x, z) = q_φ(z|x)q_φ(x) and p_θ(x, z) = p_θ(x|z)p_θ(z), we can rewrite:

= E_{z∼q_φ(z|x)}[ln (q_φ(z|x) / p_θ(x|z)p_θ(z) · p_θ(z|x)p_θ(x) / q_φ(z))]

Rearranging and simplifying:

= E_{z∼q_φ(z|x)}[ln (q_φ(z|x) / p_θ(z) · p_θ(z|x) / q_φ(z) · p_θ(x) / p_θ(x|z))]

This expands to:

I_q(x, z) = D_KL(q_φ(z|x)||p_θ(z)) − D_KL(q_φ(z)||p_θ(z)) − E_{z∼q_φ(z|x)}[ln p_θ(x|z)] + ln p_θ(x)

Noting that p_θ(z|x) ∼ p_θ(z), we obtain:

ln p_θ(x) ≥ E_{z∼q_φ(z|x)}[ln p_θ(x|z)] − D_KL(q_φ(z|x)||p_θ(z)) + D_KL(q_φ(z)||p_θ(z))

Therefore, minimizing I_q(x, z) is equivalent to maximizing the terms on the Right Hand Side (R.H.S) of the above equation, which is a lower bound.

Appendix A.3: GCVAE Constraint Optimization Proof

We recall the constraint optimization loss and prove it accordingly. Given the maximization problem:

max_{θ,φ,ξ⁺,ξ⁻,ξ_p∈R} I_p(x′, z)

s.t E_{p_D}D_KL(q_φ(z|x)||p_θ(z)) + I_p(x′, z) ≤ ξ⁻

s.t −E_{p_D}D_KL(q_φ(z)||p_θ(z)) ≤ ξ⁺

s.t I_p(x′, z) ≤ ξ_p

s.t ξ⁺_i, ξ⁻_i, ξ_ip ≥ 0, ∀i = 1, ..., n

The expansion of the above equations using sets of Lagrangian multipliers is as follows:

L(x, z; θ, φ, ξ⁺, ξ⁻, ξ, α, β, γ, η, τ, ν)

= I_p(x′, z) − β(E_{p_D}D_KL(q_φ(z|x)||p_θ(z)) + I_p(x′, z) − Σ_i=1ⁿ ξ⁻_i)

+ γ(E_{p_D}D_KL(q_φ(z)||p_θ(z)) + Σ_i=1ⁿ ξ⁺_i)

− α(I(x′; z) − Σ_i=1ⁿ ξ_ip)

+ Σ_i=1ⁿ η_iξ⁺_i + Σ_i=1ⁿ τ_iξ⁻_i + Σ_i=1ⁿ ν_iξ_ip

Simplifying and collecting terms:

L(x, z; θ, φ, ξ⁺, ξ⁻, ξ, α, β, γ, η, τ, ν)

= (1 − α − β)I_p(x′, z) − βE_{p_D}D_KL(q_φ(z|x)||p_θ(z))

+ γD_KL(q_φ(z)||p_θ(z))

+ (β − τ)Σ_i=1ⁿ ξ⁻_i + (γ − η)Σ_i=1ⁿ ξ⁺_i + (α − ν)Σ_i=1ⁿ ξ_ip

We take the gradient over the loss, ▽L for ξ⁻, ξ⁺, ξ_p and apply KKT optimality conditions to obtain:

L(x, z; θ, φ, ξ⁺, ξ⁻, ξ_p, α, β, γ)

= (1 − α − β)I_p(x′, z) − βE_{p_D}D_KL(q_φ(z|x)||p_θ(z))

+ γD_KL(q_φ(z)||p_θ(z))

Substituting I_p(x′, z) = E_{z∼q_φ(z|x)}[ln p_θ(x|z)] (since the marginal log-likelihood is intractable and can be dropped):

= (1 − α − β)E_{z∼q_φ(z|x)}[ln p_θ(x|z)]

− βE_{p_D}D_KL(q_φ(z|x)||p_θ(z))

+ γD_KL(q_φ(z)||p_θ(z))

We set the Lagrangian adaptive hyperparameters following ControlVAE as follows:

L(θ, φ, α, β, γ) = (1 − α_t − β_t)E_{z∼q_φ(z|x)}[ln p_θ(x|z)]

− β_tE_{p_D}D_KL(q_φ(z|x)||p_θ(z))

+ γ_tD_KL(q_φ(z)||p_θ(z))

The adaptive weight α_t controls the reconstruction error while β_t ensures the posterior latent factor q_φ(z|x) does not deviate significantly from its prior p_θ(z). Varying both terms gives us better control of the degree of disentanglement and helps us understand the parameters affecting density disentanglement.

Appendix A.4: Expected Squared Mahalanobis Distance

Suppose that a data x ∈ X with probability p(x) is projected into a reproducing kernel Hilbert space ⟨φ(x)⟩_H, the expectation of the squared Mahalanobis distance is as follows:

ED²_MAH(q(z)||p(z)) ≥ ||E_x∼pφ(x) − E_y∼qφ(y)||²_Σ⁻¹,H

Expanding the squared norm:

≥ (E_x∼pφ(x) − E_y∼qφ(y))′Σ⁻¹(E_x∼pφ(x) − E_y∼qφ(y))

Expanding the quadratic form:

≥ Σ⁻¹(E_x∼pφ(x) − E_y∼qφ(y))′(E_x∼pφ(x) − E_y∼qφ(y))

Expanding further:

≥ Σ⁻¹[E_{x′∼p,x∼p}⟨φ(x′), φ(x)⟩_H

− 2E_{x′∼p,y∼q}⟨φ(x′), φ(y)⟩_H

+ E_{y′∼q,y∼q}⟨φ(y′), φ(y)⟩_H]

We suppose that the feature map φ(x) takes the canonical form k(x, ·), so that ⟨φ(x′), φ(x)⟩_H = k(x′, x) where k(x′, x) represents a positive semi-definite kernel (for instance the Gaussian kernel, exp(−||x′ − x||² / 2σ²)). Hence:

ED²_MAH(q(z)||p(z)) ≥ Σ⁻¹[E_{z′∼p,z∼p}k(z′, z)

− 2E_{z′∼p,z∼q}k(z′, z)

+ E_{z′∼q,z∼q}k(z′, z)]

This simplifies to:

ED²_MAH(q(z)||p(z)) = Σ⁻¹D²_MMD(q(z)||p(z))

ED²_MAH(q||p) is a measure of the average dissimilarity between the distributions p and q in the Hilbert space, as it normalizes the similarity with feature variances, therefore encouraging class discrimination. ED²_MAH(q||p) reduces to D²_MMD(q||p) when Σ⁻¹ is identity.

Key Insights from the Derivations

The derivations establish several important theoretical results:

Equivalence of Objectives: Maximizing negative KL-divergence is equivalent to maximizing ELBO, providing a principled lower bound on the log-likelihood.
Mutual Information Minimization: Minimizing I_q(x, z) equivalently maximizes components of the ELBO, establishing the connection between information theory and variational inference.
Constraint Optimization: The GCVAE framework can be derived from first principles using Lagrangian multipliers and KKT conditions, providing a rigorous mathematical foundation.
Mahalanobis Distance: The squared Mahalanobis distance provides a better disentangling metric than Maximum Mean Discrepancy (MMD) by normalizing with feature variances.

GCVAE Loss Function

The final GCVAE loss function, derived from the constraint optimization framework, is:

L(θ, φ, α, β, γ) = (1 − α_t − β_t)E_{z∼q_φ(z|x)}[ln p_θ(x|z)]

− β_tE_{p_D}D_KL(q_φ(z|x)||p_θ(z))

+ γ_tD_KL(q_φ(z)||p_θ(z))

where:

α_t: Controls the reconstruction error (adaptive hyperparameter)
β_t: Ensures the posterior q_φ(z|x) does not deviate significantly from the prior p_θ(z)
γ_t: Controls the correlation measure (Mahalanobis distance or MMD)

The adaptive hyperparameters α_t, β_t, and γ_t are controlled using PID controllers similar to ControlVAE, allowing automatic balancing of reconstruction and disentanglement objectives.

Relationship to Other VAE Variants

The GCVAE framework generalizes several existing VAE variants. By setting specific values for the hyperparameters, we can recover:

ELBO: When α_t = α = −1, β_t = β = 1, γ = 0
ControlVAE: When α_t = α = 0, β_t > 0, γ = 0
InfoVAE: When α_t = α = 0, β_t = β = 0, γ_t > 1
FactorVAE: When α_t = α = −1, β_t = β = 1, γ = −1

This demonstrates that GCVAE provides a unified framework that encompasses and extends existing variational autoencoder approaches.

Experimental Results

We evaluate the performance of GCVAE on standard benchmark datasets including DSprites and MNIST, comparing against state-of-the-art VAE variants.

Performance Comparison

Table 1 presents a comprehensive performance comparison of different models on the DSprites dataset after training on 737 samples. The comparison metrics include MIG (Mutual Information Gap), Modularity, JEMMIG (Joint Entropy Minus Mutual Information Gap), reconstruction loss, and KL loss.

Model	MIG ↑	Modularity ↑	JEMMIG ↑	Reconstruction Loss ↓	KL Loss ↗
VAE	0.1268	0.798	0.233	3.339	3.0025
β-VAE	0.0778	0.881	0.238	0.012	35.0295
ControlVAE	0.1213	0.782	0.312	0.016	24.3809
InfoVAE	0.1501	0.757	0.188	0.079	10.0621
GCVAE-I	0.1507	0.844	0.236	0.012	24.3739
GCVAE-II	0.2793	0.858	0.312	0.012	24.4316
GCVAE-III	0.1337	0.825	0.294	0.015	24.2937

Table 1: Performance comparison of different models on DSprites after training on 737 samples. Comparison metrics MIG, Modularity and JEMMIG for 10-D Latent representation. Higher is better for MIG, Modularity and JEMMIG. GCVAE-II performs best on MIG disentanglement metric, robustness and interpretability, plus having the lowest reconstruction error.

GCVAE-II demonstrates superior performance across multiple metrics. It achieves the highest MIG score (0.2793), indicating the best disentanglement quality. The model also maintains low reconstruction error (0.012) while achieving high modularity (0.858) and JEMMIG scores (0.312). This represents an 85% increase in MIG estimation for GCVAE-II compared to GCVAE-I, demonstrating the advantage of using the squared Mahalanobis distance as a correlation measure.

Generative Process Comparison

We evaluate the quality of generation by considering the explicitness and coherency of the encoded latent variables. The generation by GCVAE-I, II and III far outperformed those of the benchmark models, especially on the MNIST dataset.

Figure 2: Generative process comparison for the different models on DSprites after training on less than 800 samples. Model GCVAE-I, II and ControlVAE clearly outperformed other models. The reconstruction error of GCVAE-II is the lowest from Table 1.

Figure 3: Generative process comparison for the different models after training the MNIST dataset for 500 epochs. GCVAE-II and ELBO (VAE) have a similar reconstruction quality with better interpretation. GCVAE-II clearly outperformed the benchmark models in generating clear and meaningful representations of the original data.

Model Performance Analysis

Figure 4 illustrates the model performance comparison on 737 samples of DSprites data, showing the relationship between reconstruction error, KL divergence, and disentanglement metrics across different latent dimensions.

Model performance comparison on DSprites

Figure 4: Model performance comparison on 737 samples of DSprites data. Top: Comparison of reconstruction error against KL divergence D_KL(q_φ(z|x)||p_θ(z)) and correlation D_KL(q_φ(z)||p_θ(z)). Bottom: Comparing disentanglement metrics with reconstruction loss. The highest disentanglement on the MIG metric is observed for GCVAE-II on Latent-2, however, the best scores are observed for GCVAE-II on Latent-10.

Stopping Criterion Analysis

We compare the performance of GCVAE-I, II & III using a stopping criterion algorithm. Using a fixed ϵ for α_t = 10e−5 and β_t = 10e−4, we obtain reasonably low reconstruction loss with high disentanglement without having to train for a lengthy period. The average time required to train GCVAE using the stopping criterion is 6 hours while it takes more than 3 days to train for 250K iterations.

Figure 5: Comparison of metrics for DSprites 2D shapes dataset. Top: Comparison of GCVAE-I, II & III losses over increasing latent dimensions. The lowest reconstruction error is observed for GCVAE-I on Latent-10 and is monotone increasing thereafter. Bottom: Disentanglement metric over different dimensions. FactorVAE metric is monotone decreasing for latent space greater than 2. MIG is similar in behavior to FactorVAE metric and best for GCVAE-II on Latent-2.

Visual comparison of reconstruction and KL losses

Figure 6: A visual comparison of the reconstruction, −D_KL(q_φ(z|x)||p_θ(z)) and D_KL(q_φ(z)||p_θ(z)) losses for GCVAE-I, II & III over different latent space. Behavior of reconstruction error per latent is relatively close and indistinguishable. In all cases of latents experimented with except for Latent-2, −D_KL(q_φ(z|x)||p_θ(z)) is comparable.

We observe that GCVAE-I is unstable in D_KL(q_φ(z)||p_θ(z)) during training across all latents. While a lower value of D_KL(q_φ(z)||p_θ(z)) is preferred, we observe correlation with MIG disentanglement metric in Figure 5. This highlights the importance of the correlation term in achieving optimal disentanglement.

Conclusion

The Generalized-Controllable Variational AutoEncoder (GCVAE) provides a principled framework for balancing reconstruction quality and latent space disentanglement. Through rigorous mathematical derivations, we establish that:

Maximizing mutual information in the reconstruction network is equivalent to information maximization during amortized inference under reasonable constraints
The constraint optimization formulation leads to a loss function with three controllable hyperparameters that automatically balance reconstruction and disentanglement
The squared Mahalanobis distance provides a superior metric for measuring disentanglement compared to Maximum Mean Discrepancy
GCVAE generalizes existing VAE variants, providing a unified framework for variational inference

The detailed derivations presented in this blog post establish the theoretical foundation of GCVAE and demonstrate its advantages over existing approaches in simultaneously achieving high disentanglement scores and low reconstruction errors.

Read the Full Paper

For additional experimental results, implementation details, and comprehensive analysis, please refer to the full research paper available on arXiv.

Access Full Paper on arXiv

Introduction

Background: Variational Autoencoders

The GCVAE Framework

Mathematical Derivations

Appendix A.1: Proof that Maximizing Negative KL-Divergence is Equivalent to Maximizing ELBO

Appendix A.2: Minimizing Mutual Information Iq(x, z)

Appendix A.3: GCVAE Constraint Optimization Proof

Appendix A.4: Expected Squared Mahalanobis Distance

Key Insights from the Derivations

GCVAE Loss Function

Relationship to Other VAE Variants

Experimental Results

Performance Comparison

Generative Process Comparison

Model Performance Analysis

Stopping Criterion Analysis

Conclusion

Read the Full Paper

Appendix A.2: Minimizing Mutual Information I_q(x, z)