Question about codebook collapse during training

[BingZhouW]

Hi, thanks for sharing this great work and the code!

I’m currently experimenting with a similar setup and had a question regarding the behavior of the codebook during training.

Specifically, I’m wondering whether you observed any codebook collapse issues (e.g., low perplexity, low active code ratio, high code reuse) when training your model. In my own experiments, I sometimes see the model relying on only a small subset of codes.

If you did encounter such issues, I would really appreciate it if you could share:

• Whether this was a problem in your setting
• What techniques you found effective to mitigate it

Thanks a lot for your time and for making your work available!

[QuLiao1117]

Hi,

Thanks for your interest!

We did not observe codebook collapse in our training. In our setting, the codebook utilization usually increased very quickly and reached 100%.

I’m not sure what modifications you made compared with our setup. Could you share more details about your changes or training configuration?

Best

[BingZhouW]

Thank you very much for your reply! I really appreciate your help.

I would like to provide more details about the issue I encountered.

My goal is to train a VQ-VAE to reconstruct the original CLIP features. I used reconstruction loss, commitment loss, and codebook loss. In my initial setting, the weights were 1.0, 0.25, and 0.5, respectively. The codebook size was 8192, and the latent dimension was 32.

I trained the model on the ImageNet-1K training set, using the last-layer CLIP features as reconstruction targets.

During training, I found that the codebook usage and perplexity were both very low, which looked like codebook collapse. I then tried reducing the codebook size to 1024 and 512, and also changed the loss weights by setting the codebook loss weight to 0.8 and the commitment loss weight to 1.0. However, the collapse issue still remained.

Later, I initialized the codebook with KMeans using the dimension-reduced CLIP features. This helped a lot with the collapse problem and significantly improved the codebook usage. However, I found that the learned code vectors became very similar to each other. The differences between different codes were quite small. Also, when I used the quantized codes for linear probing classification, the accuracy was only around 20%.

So I am wondering if this may be caused by the fact that CLIP features are already relatively close in the high-dimensional space, and after dimensionality reduction and quantization, the code vectors may become even more similar to each other. This might make the discrete codes less discriminative for downstream tasks.

Have you encountered similar behavior when reconstructing CLIP features with VQ-VAE? Do you have any suggestions for improving the diversity of the codebook or preserving the discriminative structure of CLIP features during quantization?

Thanks again for your time and help!😢😢😢😢

[QuLiao1117]

Hi,

Thanks for sharing the details.

Just to confirm, are you training a single-branch VQ model to reconstruct CLIP features? If so, codebook collapse can indeed happen quite easily in this setting.

The reason is that, at the beginning of training, the features in the codebook are not yet well separated. Meanwhile, the encoder is optimized to move its outputs closer to the selected code vectors. As a result, many patches may be assigned to the same code indices early on, causing the encoder to quickly learn nearly fixed/similar features. Once the encoder behaves this way, the codebook may also be pulled toward similar features, leading to collapse.

A simple thing to try is to reduce the learning rate of the encoder while keeping the learning rates of the other modules unchanged. Another option is to freeze the encoder at the beginning and train the other parts for a while before unfreezing it.

For further questions, feel free to add my WeChat: leo_qla

Best.

[BingZhouW]

Thank you very much for your suggestion!!!

I will try lowering the encoder learning rate and train the model again. I sincerely appreciate your reply and help.

Wish you all the best with your work and life!