| --- |
| license: mit |
| tags: |
| - vector-quantization |
| - image-tokenizer |
| - codebook-regularization |
| - icml2026 |
| datasets: |
| - imagenet-1k |
| --- |
| |
| # DimVQ: Unveiling And Addressing Dimensional Collapse In Vector Quantization Models Via Codebook Regularization |
|
|
| Official pre-trained checkpoints for the **ICML 2026** paper. |
|
|
| ## Model Description |
|
|
| DimVQ identifies **dimensional collapse** in vector quantization models and proposes a simple **codebook regularization** to restore suppressed low-variance components. This regularization bridges the spectral gap between discrete codebook spaces and continuous representations. |
|
|
| ## Available Checkpoints |
|
|
| | File | Model | Resolution | Codebook Size (K) | Embedding Dim (D) | |
| |------|-------|-----------|-------------------|-------------------| |
| | `simvq_K65536/65536.ckpt` | SimVQ + Codebook Reg. | 128x128 | 65,536 | 128 | |
| | `simvq_K65536/65536.yaml` | Config for above | - | - | - | |
| | `simvq_K262144/262144.ckpt` | SimVQ + Codebook Reg. | 128x128 | 262,144 | 128 | |
| | `simvq_K262144/262144.yaml` | Config for above | - | - | - | |
|
|
| ## Usage |
|
|
| ```python |
| # Load checkpoint |
| import torch |
| checkpoint = torch.load("262144.ckpt", map_location="cpu") |
| model.load_state_dict(checkpoint["state_dict"]) |
| ``` |
|
|
| ## TODO |
|
|
| - [ ] IBQ checkpoints (K=16384, K=262144, 256x256) |
| - [ ] Downstream autoregressive generation models (IBQ-B, IBQ-L, IBQ-XXL) |
|
|
| ## Citation |
|
|
| ```bibtex |
| @inproceedings{zhang2026dimvq, |
| title={Unveiling And Addressing Dimensional Collapse In Vector Quantization Models Via Codebook Regularization}, |
| author={Zhang, Fang and Zhu, Yongxin and Liu, Yihao and Fu, Bin and Xu, Linli}, |
| booktitle={International Conference on Machine Learning (ICML)}, |
| year={2026} |
| } |
| ``` |
|
|
| ## Links |
|
|
| - [Paper (arXiv)](https://arxiv.org/abs/TODO) |
| - [Code (GitHub)](https://github.com/ksblk2116/dimvq) |
|
|
