[Paper Review] Improving Disentangled Representation Learning with the Beta Bernoulli Process

Prashnna Gyawali, Zhiyuan Li+5 more|arXiv (Cornell University)|Sep 3, 2019

Adversarial Robustness in Machine Learning39 references4 citations

TL;DR

This paper proposes IBP-VAE, a variational autoencoder that uses the Indian Buffet Process (IBP) as a non-parametric prior to enhance disentangled representation learning. By enabling richer posterior modeling capacity while preserving factor independence, IBP-VAE improves disentanglement—especially on complex data—outperforming state-of-the-art methods on MNIST, 3D Chairs, dSprites, and clinical ECG and skin lesion datasets, with superior disentanglement of nuisance factors and improved downstream task accuracy.

ABSTRACT

To improve the ability of VAE to disentangle in the latent space, existing works mostly focus on enforcing independence among the learned latent factors. However, the ability of these models to disentangle often decreases as the complexity of the generative factors increases. In this paper, we investigate the little-explored effect of the modeling capacity of a posterior density on the disentangling ability of the VAE. We note that the independence within and the complexity of the latent density are two different properties we constrain when regularizing the posterior density: while the former promotes the disentangling ability of VAE, the latter -- if overly limited -- creates an unnecessary competition with the data reconstruction objective in VAE. Therefore, if we preserve the independence but allow richer modeling capacity in the posterior density, we will lift this competition and thereby allow improved independence and data reconstruction at the same time. We investigate this theoretical intuition with a VAE that utilizes a non-parametric latent factor model, the Indian Buffet Process (IBP), as a latent density that is able to grow with the complexity of the data. Across three widely-used benchmark data sets and two clinical data sets little explored for disentangled learning, we qualitatively and quantitatively demonstrated the improved disentangling performance of IBP-VAE over the state of the art. In the latter two clinical data sets riddled with complex factors of variations, we further demonstrated that unsupervised disentangling of nuisance factors via IBP-VAE -- when combined with a supervised objective -- can not only improve task accuracy in comparison to relevant supervised deep architectures but also facilitate knowledge discovery related to task decision-making. A shorter version of this work will appear in the ICDM 2019 conference proceedings.

Motivation & Objective

• To address the limitation of existing VAEs in disentangling complex generative factors due to constrained posterior modeling capacity.
• To investigate the trade-off between posterior independence and modeling capacity in VAEs, particularly how limited capacity undermines disentanglement.
• To improve disentanglement by increasing posterior modeling capacity while preserving factor independence through a non-parametric prior.
• To demonstrate the effectiveness of IBP-VAE in unsupervised disentanglement of nuisance factors in clinical data, enhancing interpretability and downstream task performance.
• To validate the method on benchmark datasets and underexplored clinical datasets with complex, high-dimensional variations.

Proposed method

• The IBP-VAE employs the Indian Buffet Process (IBP) as a non-parametric prior to model an unbounded number of independent latent factors, enabling flexible posterior capacity.
• The IBP prior allows the model to grow in complexity with data, avoiding the tension between reconstruction and disentanglement caused by overly constrained posteriors.
• Concrete and Kumaraswamy distributions are used to differentiably approximate Bernoulli and Beta variables in the IBP, enabling end-to-end training via backpropagation.
• The model uses a standard VAE objective with KL divergence regularization, but the prior structure ensures factor independence while allowing richer posterior representations.
• For clinical data, a conditional IBP-VAE (cIBP-VAE) is introduced to disentangle task-relevant and nuisance factors, such as pacing artifacts in ECG signals.
• The architecture uses convolutional encoders and decoders with batch normalization and ReLU activations, tailored to image and clinical signal data.

Figure 1: Outline of the presented IBP-VAE (blue) for unsupervised disentangled representation learning, and cIBP-VAE (blue and orange) for the combination with supervised task learning.

Experimental results

Research questions

• RQ1 Does increasing the modeling capacity of the posterior density in VAEs improve disentanglement, especially when generative factors are complex?
• RQ2 Can a non-parametric prior like the IBP maintain factor independence while allowing richer posterior representations to reduce conflict between reconstruction and disentanglement objectives?
• RQ3 How does IBP-VAE perform on benchmark datasets with diverse generative factors compared to state-of-the-art disentangled VAEs?
• RQ4 Can unsupervised disentanglement of nuisance factors via IBP-VAE improve supervised task performance in clinical data with high variation?
• RQ5 To what extent can IBP-VAE uncover and disentangle interpretable, semantically meaningful factors in real-world clinical data?

Key findings

• On MNIST, 3D Chairs, and dSprites, IBP-VAE achieved state-of-the-art disentanglement performance, outperforming methods like β-VAE and β-TCVAE in qualitative and quantitative evaluations.
• In the dSprites dataset, IBP-VAE successfully disentangled continuous factors such as rotation and scale, with clear disentanglement observed in disentanglement metrics.
• On the ECG dataset, cIBP-VAE reconstructed the absence of pacing artifacts significantly better than c-VAE (p < 0.01), demonstrating improved modeling of complex nuisance factors.
• cIBP-VAE successfully disentangled the pacing artifact factor, as shown by swapping representations between samples and transferring the artifact presence/absence in reconstructions.
• Deactivation of the triggering unit in cIBP-VAE introduced pacing artifacts to reconstructed signals, confirming disentanglement of this specific nuisance factor.
• In clinical skin lesion analysis, cIBP-VAE achieved comparable reconstruction accuracy to c-VAE but with superior disentanglement of artifact-related factors, enabling better interpretability and task performance.

Figure 2: [Best viewed in color] (a)-(e): Images generated by traversal along a single latent unit (over a range of [-3, 3]) on the latent representation encoded from a random sample (each row). (f): Triggering capacity of the IBP-VAE: column one: original images; column two: reconstructed images; c