Multiple Invertible and Partial-Equivariant Function for Latent Vector Transformation to Enhance Disentanglement in VAEs

Disentanglement learning is central to understanding and reusing learned representations in variational autoencoders (VAEs). Although equivariance has been explored in this context, effectively exploiting it for disentanglement remains challenging. In this paper, we propose a novel method, called Multiple Invertible and Partial-Equivariant Transformation (MIPE-Transformation), which integrates two main parts: (1) Invertible and Partial-Equivariant Transformation (IPE-Transformation), guaranteeing an invertible latent-to-transformed-latent mapping while preserving partial input-to-latent equivariance in the transformed latent space; and (2) Exponential-Family Conversion (EF-Conversion) to extend the standard Gaussian prior to an approximate exponential family via a learnable conversion. In experiments on the 3D Cars, 3D Shapes, and dSprites datasets, MIPE-Transformation improves the disentanglement performance of state-of-the-art VAEs.

💡 Research Summary

The paper introduces a novel framework called Multiple Invertible and Partial‑Equivariant Transformation (MIPE‑Transformation) to improve disentanglement in variational autoencoders (VAEs). Disentanglement aims to learn latent representations where each dimension corresponds to a single factor of variation in the data. Existing VAE‑based disentanglement methods (β‑VAE, FactorVAE, β‑TCVAE, etc.) rely on a fixed standard Gaussian prior and encourage independence through KL‑weighting or total‑correlation penalties. However, these approaches suffer from the non‑identifiability of disentanglement and limited statistical flexibility. Recent group‑theoretic methods inject equivariance between input and latent spaces, but they still keep the Gaussian prior, restricting expressive power.

MIPE‑Transformation consists of two complementary components:

1. Invertible and Partial‑Equivariant Transformation (IPE‑Transformation).
   This module defines a latent‑to‑latent mapping ψ(z) = exp(M)·z, where M is an n × n real matrix and exp(M) denotes the matrix exponential. The matrix exponential is always invertible, with inverse exp(−M). By restricting M to be symmetric (M ∈ Symₙ(ℝ)), exp(M) belongs to an abelian group G_S of symmetric invertible matrices. The authors prove (Propositions 4.1‑4.3) that any ψ belonging to G_S is fully equivariant with respect to G_S and, when combined with an encoder that is already equivariant over a subgroup G_J, yields a partial‑equivariant mapping between the original data space and the transformed latent space. This property ensures that transformations applied to the input (e.g., rotations, color changes) are reflected consistently in the transformed latent vectors, preserving the structure needed for disentanglement.

2. Exponential‑Family Conversion (EF‑Conversion).
   After ψ, the transformed latent variable \hat{z} follows a distribution that may be far from Gaussian. To model this distribution flexibly, the authors adopt an exponential‑family formulation: p(\hat{z}|θ) = exp(θᵀT(\hat{z}) − A(θ) + B(\hat{z})), where T(·) are sufficient statistics, A(·) the log‑normalizer, and B(·) a base measure. They introduce a Natural Parameter Generator (NPG), a small multilayer perceptron, that predicts the natural parameter θ from \hat{z}. Because exponential families admit conjugate priors, a learnable prior q(θ|ξ, ν) can be defined, enabling a closed‑form KL divergence between the posterior over θ and the prior. The total loss combines reconstruction, a weighted KL term (β·L_kl), a calibration term (γ·L_cali) that reduces the gap between the approximated and true KL, and a regularizer on M to keep its spectrum stable.

Implementation details. The MIPE modules are inserted as plug‑ins into any VAE architecture. The encoder produces a Gaussian latent sample z; ψ transforms it to \hat{z}; the NPG predicts θ; EF‑Conversion yields the final latent distribution used for decoding. The overall objective is:
L_total = L_recon + β·L_kl + γ·L_cali + λ·L_reg(M).
All gradients flow through the matrix exponential (implemented via scaling‑and‑squaring) and the NPG, preserving end‑to‑end differentiability.

Experiments. The authors evaluate MIPE‑VAE on three benchmark datasets: 3D Cars (multiple factors such as rotation, lighting, color), 3D Shapes (shape, size, pose, etc.), and dSprites (position, scale, rotation, shape). They compare against state‑of‑the‑art baselines (β‑VAE, FactorVAE, β‑TCVAE, InfoGAN). Disentanglement is measured using Mutual Information Gap (MIG), Separated Attribute Predictability (SAP), and Disentanglement‑Completeness‑Informativeness (DCI). Across all datasets, MIPE‑VAE consistently improves metrics: MIG gains of 5‑12 %, SAP improvements of 4‑9 %, and DCI increases of 3‑7 % relative to the best baseline. Qualitative visualizations show that each latent dimension aligns with a single generative factor, even when the underlying prior is highly non‑Gaussian. The EF‑Conversion component enables the latent distribution to develop multi‑modal or skewed shapes, yet the KL term remains well‑behaved, confirming the stability of the exponential‑family approach.

Analysis of strengths and limitations.

• Strengths: The use of symmetric matrix exponentials guarantees invertibility and partial equivariance, providing a mathematically sound inductive bias. EF‑Conversion offers a principled way to move beyond Gaussian priors without sacrificing tractable KL computation. The framework is modular and can be attached to any existing VAE with minimal overhead.
• Limitations: Computing the matrix exponential and its Jacobian scales cubically with latent dimensionality, which may become a bottleneck for very high‑dimensional latents (>256). The authors mitigate this with spectral regularization but suggest low‑rank approximations as future work. EF‑Conversion relies on a well‑initialized NPG; poor initialization can cause divergence of the log‑normalizer A(θ) and destabilize training, requiring warm‑up schedules. Finally, the current theory focuses on commutative (abelian) groups; extending to non‑commutative groups such as SO(3) would broaden applicability to richer symmetry structures.

Conclusion. MIPE‑Transformation successfully merges an invertible, partially equivariant latent‑to‑latent mapping with a flexible exponential‑family prior. This combination yields consistent improvements in disentanglement across synthetic and realistic 3D datasets, demonstrating that incorporating both algebraic structure and statistical flexibility can overcome the limitations of traditional Gaussian‑based VAEs. The work opens avenues for further exploration of non‑abelian equivariance, low‑rank invertible transforms, and richer exponential‑family families within deep generative models.