Disentanglement via Latent Quantization

Abstract: In disentangled representation learning, a model is asked to tease apart a dataset's underlying sources of variation and represent them independently of one another. Since the model is provided with no ground truth information about these sources, inductive biases take a paramount role in enabling disentanglement. In this work, we construct an inductive bias towards encoding to and decoding from an organized latent space. Concretely, we do this by (i) quantizing the latent space into discrete code vectors with a separate learnable scalar codebook per dimension and (ii) applying strong model regularization via an unusually high weight decay. Intuitively, the latent space design forces the encoder to combinatorially construct codes from a small number of distinct scalar values, which in turn enables the decoder to assign a consistent meaning to each value. Regularization then serves to drive the model towards this parsimonious strategy. We demonstrate the broad applicability of this approach by adding it to both basic data-reconstructing (vanilla autoencoder) and latent-reconstructing (InfoGAN) generative models. For reliable evaluation, we also propose InfoMEC, a new set of metrics for disentanglement that is cohesively grounded in information theory and fixes well-established shortcomings in previous metrics. Together with regularization, latent quantization dramatically improves the modularity and explicitness of learned representations on a representative suite of benchmark datasets. In particular, our quantized-latent autoencoder (QLAE) consistently outperforms strong methods from prior work in these key disentanglement properties without compromising data reconstruction.

• The paper proposes latent quantization, using separate scalar codebooks per dimension to boost disentanglement without hindering data reconstruction.
• It employs strong weight decay regularization and InfoMEC metrics to rigorously measure modularity, explicitness, and compactness.
• Experiments on datasets like Shapes3D and MPI3D demonstrate that QLAE outperforms β-VAE and β-TCVAE in achieving superior interpretability and accuracy.

Disentanglement via Latent Quantization

Introduction

The paper introduces a technique for disentangled representation learning that involves quantizing latent space into discrete code vectors with separate learnable scalar codebooks per dimension. This strategy enhances the modularity and explicitness of learned representations without compromising data reconstruction. The paper also proposes InfoMEC, a set of information-theoretic metrics to evaluate disentanglement, correcting deficiencies in former metrics.

Methodology

Latent Quantization

The authors propose a method to encourage disentanglement by organizing the latent space through quantization. This involves representing latent codes as Cartesian products of scalar codebooks, allowing efficient nearest neighbor calculations. Each latent dimension possesses a distinct codebook that regularizes the latent representation structure.

Figure 1: Depiction of how latent quantization uses combinatorial scalar codebooks to enable efficient recovery of source data.

Regularization

High weight decay is applied to both encoder and decoder networks as a regularization strategy. The method biases models towards using a more parsimonious representation, easing the recovery of true data-generating factors.

Experiments

Experiments conducted on datasets like Shapes3D, MPI3D, Falcor3D, and Isaac3D show that latent quantization dramatically improves disentanglement. QLAE outperforms traditional methods like $\beta$-VAE and $\beta$-TCVAE in modularity without sacrificing data reconstruction. Different ablations confirm the criticality of dimension-specific codebooks and model regularization.

Figure 2: Comparative analysis showing superior modularity and explicitness metrics of QLAE across various datasets.

InfoMEC Metrics

The paper presents InfoMEC, a set of metrics for disentanglement evaluation. InfoMEC includes:

• InfoModularity (InfoM): The extent to which sources are encoded into disjoint sets of latents.
• InfoExplicitness (InfoE): Simplicity of the relationship between sources and latents, measured by linear predictiveness.
• InfoCompactness (InfoC): Latents containing information disjoint from other sources.

These metrics are more robust than existing ones due to their information-theoretic foundations.

Figure 3: Visualization of NMI heatmaps for model evaluations demonstrating superior InfoMEC scores for QLAE.

Discussion

Latent quantization successfully arranges discrete latents to mimic organized source space, aligning well with realistic datagenerating models. Despite promising results, the technique faces challenges such as potentially overfitting the discrete assumptions of contemporary benchmarks. Further exploration into continuous or noisier data settings might enhance generalization.

Conclusion

The integration of latent quantization and strong regularization offers a formidable improvement for disentangled representation learning. The introduction of InfoMEC presents a robust and cohesive framework for evaluating model performance in disentanglement tasks, while the QLAE model sets a new benchmark for interpretability and accuracy. Future work could explore combining latent quantization with other structural assumptions in generative processes.

Figure 4: Decoded interventions showing the interpretability and consistency in source changes across latent dimensions.