search

IDSIA / kohonen-vae

Official repository for the paper "Topological Neural Discrete Representation Learning à la Kohonen" (ICML 2023 Workshop on Sampling and Optimization in Discrete Space)
8 stars 1 forks source link

HardSOM update method: self._weight to be replaced by self._w #1

Closed BeusenBart closed 1 year ago

BeusenBart commented 1 year ago

It seems that in som_vector_quantizer.py, in class HardSOM, in the method "update", the embedding weights are not updated correctly: on line 202: "self._weight = self._ema_w / self._ema_cluster_size.unsqueeze(1)" should be replaced by " self._w = self._ema_w / self._ema_cluster_size.unsqueeze(1)". The variable "self._weight" is used nowhere else in the code, I expect this should be self._w. After fixing this bug, I'm able to get good quality reconstructions using my own VQ-VAE implementation that uses the HardSOM as a quantizer layer. Without the fix, the model underperforms drastically. Thank you for sharing this code, I believe the Kohonen SOM is a valuable alternative for the EMA updating.

RobertCsordas commented 1 year ago

Hi! Thank you for your interest! We are very sorry for this bug, and thank you for reporting it. It was introduced when refactoring the code for public release. The experiments in the paper were done with the correct code. I pushed the fix and also a new, gradient-descent-based baseline.