Self-Organising Neural Discrete Representation Learning \`a la Kohonen

TL;DR

This paper explores an alternative vector quantisation method based on Kohonen's Self-Organising Map for neural networks, demonstrating faster convergence and greater robustness in learning discrete representations, especially in image processing tasks.

Contribution

It introduces KSOM as a robust and potentially faster alternative to EMA-VQ for discrete representation learning in neural networks, with applications to image processing.

Findings

KSOM converges faster initially than EMA-VQ.

KSOM provides a topological structure in discrete representations.

KSOM shows greater robustness to initialization schemes.

Abstract

Unsupervised learning of discrete representations in neural networks (NNs) from continuous ones is essential for many modern applications. Vector Quantisation (VQ) has become popular for this, in particular in the context of generative models, such as Variational Auto-Encoders (VAEs), where the exponential moving average-based VQ (EMA-VQ) algorithm is often used. Here, we study an alternative VQ algorithm based on Kohonen's learning rule for the Self-Organising Map (KSOM; 1982). EMA-VQ is a special case of KSOM. KSOM is known to offer two potential benefits: empirically, it converges faster than EMA-VQ, and KSOM-generated discrete representations form a topological structure on the grid whose nodes are the discrete symbols, resulting in an artificial version of the brain's topographic map. We revisit these properties by using KSOM in VQ-VAEs for image processing. In our experiments, the speed-up compared to well-configured EMA-VQ is only observable at the beginning of training, but KSOM is generally much more robust, e.g., w.r.t. the choice of initialisation schemes.