# Dimensional Collapse in VQ-VAEs: Evidence and Remedies

## Requirements
Create a conda environment with `Python >= 3.8` using the `environment.yml` file.

## Datasets
CelebA and CIFAR-10 will be automatically downloaded when the code is executed.

## Scripts

### Figures 2 & 3: Pretrained Model Analysis
Run `pretrain_figures/pretrained_effdim.ipynb` to reproduce Figure 2 and 3.

### Figure 4-6: 512 W&B Sweeps
1. Run:
```bash
wandb sweep src/configs/sweep_<config_name>.yaml
```
2. Then launch runs:
```bash
SLURM_PROCID=<gpu_id> CUDA_VISIBLE_DEVICES=<gpu_id> wandb agent <sweep_id> --count 1
```
Repeat 64 times per config to match the 512 runs in the paper.

3. After finishing the sweeps, run the notebooks in `512_figures/` to reproduce Figure 4-6.

### Figure 7: Koleo Experiments

1. Run `src/configs/gen_all_configs_koleo.py` to generate sweep scripts.
2. Then, execute:
```
bash src/configs/all_configs_koleo/config_<node_id>_<gpu_id>.sh
```
Adjust GPU settings as needed.

3. After finishing all the commands, run the notebooks in `koleo_figures/` to reproduce Figure 7.

### Figure 9: DCVQ Experiments

1. Run `src/configs/gen_all_configs.py` to generate sweep scripts.
2. Then, execute:
```
bash src/configs/all_configs/config_<node_id>_<gpu_id>.sh
```
Adjust GPU settings as needed.

*Note: The model referred to as DCVQ in the paper is named split_vq in the code.*

3. After finishing all the commands, run the notebooks in `dcvq_figures/` to reproduce Figure 9.

## Notes

Before starting each experiment, it's recommended to backup and clear the `rot_vqvae` directory to avoid overwriting previous runs.

## Acknowledgments

This codebase is based on [cfifty/rotation_trick](https://github.com/cfifty/rotation_trick), originally released with the ICLR 2024 paper *"Restructuring Vector Quantization with the Rotation Trick"*. We thank the authors for releasing their code under academic fair use.
