This repository implements a novel disentanglement framework that integrates Mutual Information Neural Estimation (MINE) with cross-attention mechanisms for dynamic latent capacity allocation, as described in our research paper.
Note: This repository is based on the PyTorch-VAE framework.
Disentangling data into its compositional generative factors has been a long-standing challenge in deep representation learning. Existing variational approaches often struggle with disentangling the latent space to generative factors in high-dimensional latent spaces. Previous works rely on empirical estimates of mutual information, such as total correlation, which require Monte Carlo sampling and suffer from high variance, limiting their effectiveness in complex data distributions.
To address these limitations, we propose a novel disentanglement framework integrating the mutual information neural estimator (MINE) with a cross-attention mechanism for dynamic capacity allocation in the latent space. Our approach promotes disentanglement and adaptively assigns latent capacity for each generative factor based on its complexity. We allow high-dimensional partitions for each generative factor to capture complex factors while preserving some of the latent code capacity for nuisance factors that are free to remain entangled.
- Outer Loop: Optimizes disentanglement objectives using MINE estimators
- Inner Loop: Adapts cross-attention weights for dynamic capacity allocation
- Convergence Guarantees: Theoretical analysis of optimization convergence
- Mutual Information Neural Estimator: Implements the MINE framework for accurate mutual information estimation
- Low Variance Estimation: Reduces estimation bias compared to Monte Carlo methods
- Neural Network-based: Uses neural networks for flexible mutual information estimation
- Dynamic Capacity Allocation: Adaptively assigns latent capacity based on factor complexity
- Soft Associations: Learns soft associations between latent codes and generative factors
- Compact Embeddings: Ensures efficient use of latent space
- Factor Complexity Estimation: Measures the complexity of each generative factor
- Adaptive Weighting: Dynamically adjusts disentanglement weights based on factor complexity
- Robust Estimation: Handles varying factor distributions effectively
- MutualInformationEstimator: Neural network-based mutual information estimation
- CrossAttention: Dynamic capacity allocation mechanism
- FactorEncoder: Encodes input to factor embeddings
- GaussianFitEstimator: Estimates factor complexity using Gaussian fitting
- MINEDisentangleVAE: Main model integrating all components
Input Image → Encoder → Factor Encoder → Cross Attention → Decoder → Output
↓ ↓ ↓
Latent Code Factor Emb. Attended Latent
↓ ↓ ↓
MINE Estimators ← Mutual Information Estimation
↓
Gaussian Fit Estimators ← Factor Complexity
# (Recommended) Create a virtual environment first
python -m venv mine_vae_env
source mine_vae_env/bin/activate
# Install all dependencies
pip install -r requirements.txt- torch==2.1.2
- torchvision==0.16.2
- torchaudio==2.1.2
- pytorch-lightning==2.5.2
- h5py>=3.10.0
- numpy>=1.21.0
- scipy>=1.7.0
- scikit-learn>=1.0.0
- matplotlib>=3.5.0
- seaborn>=0.11.0
- tensorboard>=2.8.0
- lpips>=0.1.4
- tqdm>=4.62.0
- Pillow>=8.3.0
If you prefer conda, create an environment and use pip for PyTorch Lightning:
conda create -n mine_vae_env python=3.9
conda activate mine_vae_env
pip install -r requirements.txt# Train the MINE Disentangle VAE
python run_enhanced.py -c configs/mine_disentangle_vae.yaml -e disentanglement# Evaluate a trained model
python run_enhanced.py -c configs/mine_disentangle_vae.yaml -e disentanglement -eval -cp path/to/checkpoint.ckpt# Run ablation study
python run_enhanced.py --ablationmodel_params:
name: 'MINEDisentangleVAE'
in_channels: 3
latent_dim: 128
num_factors: 10
factor_dim: 32
attention_dim: 128
mine_hidden_dim: 128
disentanglement_weight: 1.0
attention_reg_weight: 0.1
entropy_reg_weight: 0.05
bi_level_outer_steps: 5
bi_level_inner_steps: 3
gaussian_fit_samples: 1000latent_dim: Total dimensionality of the latent spacenum_factors: Number of generative factors to disentanglefactor_dim: Dimensionality of each factor embeddingdisentanglement_weight: Weight for disentanglement lossattention_reg_weight: Weight for L2 regularization on attention weightsentropy_reg_weight: Weight for entropy regularization on attention weightsbi_level_outer_steps: Number of outer optimization stepsbi_level_inner_steps: Number of inner optimization steps
- Z-diff: Measures sensitivity of latent variables to changes in generative factors (higher is better)
- Z-min Var: Measures variance and separation across dimensions (lower is better)
- IRS: Interventional Robustness Score - quantifies how well each latent dimension uniquely corresponds to a generative factor (higher is better)
- MIG: Mutual Information Gap - measures the degree of exclusivity between latent dimensions and factors (higher is better)
- JEMMIG: Joint Entropy Minus Information Gap - extends MIG by incorporating joint entropy (higher is better)
- Latent Traversals: Visual evaluation of disentanglement by sweeping across latent variables and decoding samples
3DShapes:
- Download 3dshapes.h5
- Place in:
Data/3dshapes/3dshapes.h5
dSprites:
- Download dsprites_ndarray_co1sh3sc6or40x32y32_64x64.npz
- Place in:
Data/dsprites/dsprites_ndarray_co1sh3sc6or40x32y32_64x64.npz
MPI3D:
- Download and unzip mpi3d_realistic.zip
- Place in:
Data/mpi3d/mpi3d_realistic/
CelebA:
- Download and unzip img_align_celeba.zip
- Place images in:
Data/celeba/img_align_celeba/ - Place label files (e.g.,
list_attr_celeba.txt,identity_CelebA.txt,list_bbox_celeba.txt) in:Data/celeba/
| Dataset | Download Location in Repo | File(s) Needed | Labels Included? |
|---|---|---|---|
| 3DShapes | Data/3dshapes/3dshapes.h5 | 3dshapes.h5 | Yes |
| dSprites | Data/dsprites/dsprites_ndarray_...npz | dsprites_ndarray_co1sh3sc6or40x32y32_64x64.npz | Yes |
| MPI3D | Data/mpi3d/mpi3d_realistic/ | .npz files in extracted folder | Yes |
| CelebA | Data/celeba/img_align_celeba/ | image files, .txt label files | Yes |